Feature: Reworked record combiner to support redefines in the writer

cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParametersParser.scala

@@ -147,6 +147,7 @@ object CobolParametersParser extends Logging {
   val PARAM_WRITE_NULL_STRINGS_AS_SPACES           = "write_null_strings_as_spaces"
   val PARAM_WRITE_NULL_DISPLAY_NUMBERS_AS_ZEROS    = "write_null_display_numbers_as_zeros"
   val PARAM_WRITE_NULL_COMP3_NUMBERS_AS_ZEROS      = "write_null_comp3_numbers_as_zeros"
+  val PARAM_WRITE_STRICT_REDEFINES                 = "write_strict_redefines"
 
   val MIN_RECORDS_FOR_INDEXES = 100000
 
@@ -374,7 +375,8 @@ object CobolParametersParser extends Logging {
       isEbcdic = isEbcdic,
       nullStringsAsSpaces = parameters.getOrElse(PARAM_WRITE_NULL_STRINGS_AS_SPACES, "false").toBoolean,
       nullDisplayNumbersAsZeros = parameters.getOrElse(PARAM_WRITE_NULL_DISPLAY_NUMBERS_AS_ZEROS, "false").toBoolean,
-      nullComp3NumbersAsZeros = parameters.getOrElse(PARAM_WRITE_NULL_COMP3_NUMBERS_AS_ZEROS, "false").toBoolean
+      nullComp3NumbersAsZeros = parameters.getOrElse(PARAM_WRITE_NULL_COMP3_NUMBERS_AS_ZEROS, "false").toBoolean,
+      strictRedefines = parameters.getOrElse(PARAM_WRITE_STRICT_REDEFINES, "false").toBoolean
     )
   }
 

cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/WriterParameters.scala

@@ -20,5 +20,6 @@ case class WriterParameters(
                              isEbcdic: Boolean = true,
                              nullStringsAsSpaces: Boolean = false,
                              nullDisplayNumbersAsZeros: Boolean = false,
-                             nullComp3NumbersAsZeros: Boolean = false
+                             nullComp3NumbersAsZeros: Boolean = false,
+                             strictRedefines: Boolean = false
                            )

cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParametersParserSuite.scala

@@ -76,14 +76,16 @@ class CobolParametersParserSuite extends AnyWordSpec {
         "write_null_strings_as_spaces" -> "false",
         "write_null_display_numbers_as_zeros" -> "true",
         "write_null_comp3_numbers_as_zeros" -> "true",
+        "write_strict_redefines" -> "true",
         "pedantic" -> "true"
       ))
 
       val parsedParams = CobolParametersParser.parse(params, isWriter = true)
       assert(parsedParams.writerParameters.get == WriterParameters(
         nullStringsAsSpaces = false,
         nullDisplayNumbersAsZeros = true,
-        nullComp3NumbersAsZeros = true
+        nullComp3NumbersAsZeros = true,
+        strictRedefines = true
       ))
     }
   }

spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/writer/NestedRecordCombiner.scala

@@ -22,7 +22,7 @@ import org.apache.spark.sql.{DataFrame, Row}
 import org.slf4j.LoggerFactory
 import za.co.absa.cobrix.cobol.parser.Copybook
 import za.co.absa.cobrix.cobol.parser.ast.datatype.{AlphaNumeric, COMP3, Decimal, Integral}
-import za.co.absa.cobrix.cobol.parser.ast.{Group, Primitive}
+import za.co.absa.cobrix.cobol.parser.ast.{Group, Primitive, Statement}
 import za.co.absa.cobrix.cobol.parser.policies.VariableSizeOccursPolicy
 import za.co.absa.cobrix.cobol.parser.recordformats.RecordFormat
 import za.co.absa.cobrix.cobol.reader.parameters.{ReaderParameters, WriterParameters}
@@ -234,18 +234,72 @@ object NestedRecordCombiner {
     * @param path         The path to the field
     * @param dependeeMap  A map of field names to their corresponding DependingOnField specs, used to resolve dependencies for OCCURS DEPENDING ON fields.
     * @param strictSchema If true, each field in the copybook must exist in the Spark schema.
-    * @return A [[GroupField]] covering all non-filler, non-redefines children found in both
-    *         the copybook and the Spark schema.
+    * @return A [[GroupField]] covering all children found in both the copybook and the Spark
+    *         schema. Fields that participate in a REDEFINES chain are grouped together into a
+    *         single [[RedefineGroup]] node representing all mutually exclusive alternatives.
     */
   private def buildGroupField(group: Group, schema: StructType, getter: GroupGetter, path: String, dependeeMap: mutable.HashMap[String, DependingOnField], strictSchema: Boolean): GroupField = {
-    val children = group.children.withFilter { stmt =>
-      stmt.redefines.isEmpty
-    }.map {
-      case s if s.isFiller => Filler(s.binaryProperties.actualSize)
-      case p: Primitive    => buildPrimitiveNode(p, schema, path, dependeeMap, strictSchema)
-      case g: Group        => buildGroupNode(g, schema, path, dependeeMap, strictSchema)
+    val rawChildren = group.children
+    val processed = new mutable.ArrayBuffer[WriterAst]()
+
+    var i = 0
+    while (i < rawChildren.length) {
+      val stmt = rawChildren(i)
+      // A REDEFINES chain starts at a non-redefining field and is immediately followed
+      // (in declaration order) by one or more fields that redefine an earlier field of the chain.
+      // This mirrors the clustering logic used by BinaryPropertiesAdder when computing binary sizes.
+      var j = i + 1
+      while (j < rawChildren.length && rawChildren(j).redefines.nonEmpty) {
+        j += 1
+      }
+      val clusterStmts = rawChildren.slice(i, j)
+
+      if (clusterStmts.length == 1) {
+        processed += buildChildNode(stmt, schema, path, dependeeMap, strictSchema)
+      } else {
+        processed += buildRedefineGroup(clusterStmts.toSeq, schema, path, dependeeMap, strictSchema)
+      }
+
+      i = j
     }
-    GroupField(children.toSeq, group, getter)
+    GroupField(processed.toSeq, group, getter)
+  }
+
+  /**
+    * Builds a single [[WriterAst]] node for a copybook statement, dispatching to the
+    * appropriate builder based on whether the statement is a filler, primitive or group.
+    */
+  private def buildChildNode(stmt: Statement, schema: StructType, path: String, dependeeMap: mutable.HashMap[String, DependingOnField], strictSchema: Boolean): WriterAst = stmt match {
+    case s if s.isFiller => Filler(s.binaryProperties.actualSize)
+    case p: Primitive    => buildPrimitiveNode(p, schema, path, dependeeMap, strictSchema)
+    case g: Group         => buildGroupNode(g, schema, path, dependeeMap, strictSchema)
+  }
+
+  /**
+    * Builds a [[RedefineGroup]] node from a chain of mutually exclusive copybook statements
+    * (a base field followed by one or more fields that REDEFINE it, directly or transitively).
+    *
+    * Individual alternatives are built without enforcing `strictSchema` since it is expected
+    * that only one alternative is present in the Spark schema for a given row; the strict
+    * check is instead performed once, at the level of the whole chain: if none of the
+    * alternatives are found in the schema, the usual strict/non-strict schema behavior applies.
+    */
+  private def buildRedefineGroup(clusterStmts: Seq[Statement], schema: StructType, path: String, dependeeMap: mutable.HashMap[String, DependingOnField], strictSchema: Boolean): RedefineGroup = {
+    val alternatives = clusterStmts.map { s =>
+      RedefineAlternative(s.name, buildChildNode(s, schema, path, dependeeMap, strictSchema = false))
+    }
+
+    val isPresent = alternatives.exists(alt => !alt.ast.isInstanceOf[Filler])
+    if (!isPresent) {
+      val fieldNames = clusterStmts.map(_.name).mkString("', '")
+      if (strictSchema) {
+        throw new IllegalArgumentException(s"None of the REDEFINES alternatives ('$fieldNames') at '$path${clusterStmts.head.name}' are found in Spark schema.")
+      } else {
+        log.warn(s"None of the REDEFINES alternatives ('$fieldNames') at '$path${clusterStmts.head.name}' are found in Spark schema. Will be replaced by filler.")
+      }
+    }
+
+    RedefineGroup(alternatives, clusterStmts.head.binaryProperties.actualSize)
   }
 
   /**
@@ -484,6 +538,49 @@ object NestedRecordCombiner {
           )
           if (variableLengthOccurs) 0 else cobolField.binaryProperties.actualSize
         }
+
+      // ── REDEFINES group (mutually exclusive alternatives sharing the same bytes) ─────
+      case RedefineGroup(alternatives, actualSize) =>
+        val populated = alternatives.filter(alt => isPopulated(alt.ast, row))
+        populated match {
+          case Seq() =>
+            // No alternative has a value for this row: leave the shared bytes as zeroes.
+            actualSize
+          case Seq(only) =>
+            writeToBytes(only.ast, row, ar, currentOffset, variableLengthOccurs, writerParameters)
+            actualSize
+          case multiple =>
+            val fieldNames = multiple.map(_.fieldName).mkString("', '")
+            if (writerParameters.strictRedefines) {
+              throw new IllegalArgumentException(
+                s"Conflicting REDEFINES fields populated on the same row: '$fieldNames'. " +
+                  s"Only one field of a REDEFINES group can have a non-null value at a time."
+              )
+            } else {
+              val chosen = multiple.head
+              log.warn(
+                s"Conflicting REDEFINES fields populated on the same row: '$fieldNames'. " +
+                  s"Writing the first populated alternative ('${chosen.fieldName}') and ignoring the rest. " +
+                  s"Set 'write_strict_redefines' to 'true' to fail instead."
+              )
+              writeToBytes(chosen.ast, row, ar, currentOffset, variableLengthOccurs, writerParameters)
+            }
+            actualSize
+        }
     }
   }
+
+  /**
+    * Determines whether a writer AST node has a non-null value to write for the given row.
+    * Used to detect which alternative(s) of a REDEFINES chain are populated for a row.
+    */
+  private def isPopulated(ast: WriterAst, row: Row): Boolean = ast match {
+    case Filler(_)                       => false
+    case PrimitiveField(_, getter)       => getter(row) != null
+    case PrimitiveDependeeField(_)       => false
+    case GroupField(_, _, getter)        => getter(row) != null
+    case PrimitiveArray(_, arrayGetter, _) => arrayGetter(row) != null
+    case GroupArray(_, _, arrayGetter, _)  => arrayGetter(row) != null
+    case RedefineGroup(alternatives, _)  => alternatives.exists(alt => isPopulated(alt.ast, row))
+  }
 }

spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/writer/WriterAst.scala

@@ -42,6 +42,8 @@ sealed trait WriterAst
   * - GroupField represents a COBOL group containing child fields with its getter function
   * - PrimitiveArray represents an array of primitive values with optional depending-on semantics
   * - GroupArray represents an array of group structures with optional depending-on semantics
+  * - RedefineGroup represents a set of mutually exclusive REDEFINES alternatives sharing the
+  *   same byte region; at most one alternative may be populated in a given row
   *
   * The depending-on fields support COBOL's OCCURS DEPENDING ON clause, where the actual number
   * of array elements is determined by the value of another field at runtime.
@@ -57,4 +59,24 @@ object WriterAst {
   case class GroupField(children: Seq[WriterAst], cobolField: Group, getter: GroupGetter) extends WriterAst
   case class PrimitiveArray(cobolField: Primitive, arrayGetter: ArrayGetter, dependingOn: Option[DependingOnField]) extends WriterAst
   case class GroupArray(groupField: GroupField, cobolField: Group, arrayGetter: ArrayGetter, dependingOn: Option[DependingOnField]) extends WriterAst
+
+  /**
+    * One alternative of a REDEFINES chain, keeping the original copybook field name for
+    * error reporting purposes alongside the constructed writer AST node for that alternative.
+    */
+  case class RedefineAlternative(fieldName: String, ast: WriterAst)
+
+  /**
+    * Represents a group of mutually exclusive fields (or groups) that occupy the same byte
+    * region of a record because one REDEFINES another (directly or transitively).
+    *
+    * At write time, at most one alternative is expected to carry a non-null value for a given
+    * row. If none carry a value, the shared bytes are left as zeroes (like a filler). If more
+    * than one carry a value, writing fails fast since it would be ambiguous which value should
+    * be encoded into the shared bytes.
+    *
+    * @param alternatives The list of mutually exclusive alternatives sharing the byte region.
+    * @param actualSize   The size, in bytes, of the shared byte region (uniform across all alternatives).
+    */
+  case class RedefineGroup(alternatives: Seq[RedefineAlternative], actualSize: Int) extends WriterAst
 }