diff --git a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParametersParser.scala b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParametersParser.scala index d48332a4..1436f85d 100644 --- a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParametersParser.scala +++ b/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 ) } diff --git a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/WriterParameters.scala b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/WriterParameters.scala index f5b1ffbc..424116ce 100644 --- a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/WriterParameters.scala +++ b/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 ) diff --git a/cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParametersParserSuite.scala b/cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParametersParserSuite.scala index 265ae423..417dc70f 100644 --- a/cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParametersParserSuite.scala +++ b/cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParametersParserSuite.scala @@ -76,6 +76,7 @@ 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" )) @@ -83,7 +84,8 @@ class CobolParametersParserSuite extends AnyWordSpec { assert(parsedParams.writerParameters.get == WriterParameters( nullStringsAsSpaces = false, nullDisplayNumbersAsZeros = true, - nullComp3NumbersAsZeros = true + nullComp3NumbersAsZeros = true, + strictRedefines = true )) } } diff --git a/spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/writer/NestedRecordCombiner.scala b/spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/writer/NestedRecordCombiner.scala index 4212752b..88ba1c88 100644 --- a/spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/writer/NestedRecordCombiner.scala +++ b/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)) + } } diff --git a/spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/writer/WriterAst.scala b/spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/writer/WriterAst.scala index fab1c01a..fa95a0fb 100644 --- a/spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/writer/WriterAst.scala +++ b/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 } diff --git a/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/writer/FixedLengthEbcdicWriterSuite.scala b/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/writer/FixedLengthEbcdicWriterSuite.scala index 3673ad4c..ec6e66ae 100644 --- a/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/writer/FixedLengthEbcdicWriterSuite.scala +++ b/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/writer/FixedLengthEbcdicWriterSuite.scala @@ -24,6 +24,8 @@ import za.co.absa.cobrix.spark.cobol.source.base.SparkTestBase import za.co.absa.cobrix.spark.cobol.source.fixtures.{BinaryFileFixture, TextComparisonFixture} import za.co.absa.cobrix.spark.cobol.utils.SparkUtils +import scala.annotation.tailrec + class FixedLengthEbcdicWriterSuite extends AnyWordSpec with SparkTestBase with BinaryFileFixture with TextComparisonFixture { import spark.implicits._ @@ -495,6 +497,406 @@ class FixedLengthEbcdicWriterSuite extends AnyWordSpec with SparkTestBase with B assert(fs.exists(path), "Output directory should exist") } } + + "write data frames using REDEFINES fields" should { + val copybookContentsWithRedefines = + """ 01 RECORD. + 05 A PIC X(1). + 05 B PIC 9(5). + 05 B1 PIC X(5) REDEFINES B. + """ + + "write using only the base field of a REDEFINES group" in { + withTempDirectory("cobol_writer_redefines") { tempDir => + val df = List(("X", 12345)).toDF("A", "B") + + val path = new Path(tempDir, "writer_redefines_base") + + df.coalesce(1) + .write + .format("cobol") + .mode(SaveMode.Overwrite) + .option("copybook_contents", copybookContentsWithRedefines) + .save(path.toString) + + val bytes = readPartFileBytes(path) + + // Expected EBCDIC data: A='X', B=12345 (DISPLAY digits) + val expected = Array[Byte]( + 0xE7.toByte, + 0xF1.toByte, 0xF2.toByte, 0xF3.toByte, 0xF4.toByte, 0xF5.toByte + ) + + assertArraysEqual(bytes, expected) + } + } + + "write using only the redefining field of a REDEFINES group" in { + withTempDirectory("cobol_writer_redefines") { tempDir => + val df = List(("X", "ABCDE")).toDF("A", "B1") + + val path = new Path(tempDir, "writer_redefines_alt") + + df.coalesce(1) + .write + .format("cobol") + .mode(SaveMode.Overwrite) + .option("copybook_contents", copybookContentsWithRedefines) + .save(path.toString) + + val bytes = readPartFileBytes(path) + + // Expected EBCDIC data: A='X', B1="ABCDE" + val expected = Array[Byte]( + 0xE7.toByte, + 0xC1.toByte, 0xC2.toByte, 0xC3.toByte, 0xC4.toByte, 0xC5.toByte + ) + + assertArraysEqual(bytes, expected) + } + } + + "fail fast when both the base and the redefining fields are populated on the same row" in { + withTempDirectory("cobol_writer_redefines") { tempDir => + val df = List(("X", 12345, "ABCDE")).toDF("A", "B", "B1") + + val path = new Path(tempDir, "writer_redefines_conflict") + + val thrown = intercept[Throwable] { + df.coalesce(1) + .write + .format("cobol") + .mode(SaveMode.Overwrite) + .option("copybook_contents", copybookContentsWithRedefines) + .option("write_strict_redefines", "true") + .save(path.toString) + } + + val messages = causeChainMessages(thrown) + assert(messages.exists(m => m.contains("B") && m.contains("B1")), + s"Expected an error mentioning both conflicting REDEFINES fields 'B' and 'B1', but got: ${messages.mkString(" | ")}") + } + } + + "write the first alternative when multiple REDEFINES fields are populated and strict is disabled" in { + withTempDirectory("cobol_writer_redefines") { tempDir => + val df = List(("X", 12345, "ABCDE")).toDF("A", "B", "B1") + + val path = new Path(tempDir, "writer_redefines_first_wins") + + df.coalesce(1) + .write + .format("cobol") + .mode(SaveMode.Overwrite) + .option("copybook_contents", copybookContentsWithRedefines) + .save(path.toString) + + val bytes = readPartFileBytes(path) + + // A='X' (0xE7), then B=12345 as EBCDIC DISPLAY digits (0xF1..0xF5). B1 ("ABCDE") is ignored. + val expected = Array[Byte]( + 0xE7.toByte, + 0xF1.toByte, 0xF2.toByte, 0xF3.toByte, 0xF4.toByte, 0xF5.toByte + ) + + assertArraysEqual(bytes, expected) + } + } + + "write a data frame with a mix of rows using the base field and rows using the redefining field" in { + withTempDirectory("cobol_writer_redefines") { tempDir => + val df = List( + ("1", Some(11111), None), + ("2", None, Some("AAAAA")), + ("3", Some(33333), None), + ("4", None, Some("BBBBB")) + ).toDF("A", "B", "B1") + + val path = new Path(tempDir, "writer_redefines_mixed_rows") + + df.coalesce(1) + .write + .format("cobol") + .mode(SaveMode.Overwrite) + .option("copybook_contents", copybookContentsWithRedefines) + .save(path.toString) + + val bytes = readPartFileBytes(path) + + // Expected EBCDIC data per row: A=, then either B (DISPLAY digits) or B1 (alpha) depending on which is populated + val expected = Array[Byte]( + 0xF1.toByte, 0xF1.toByte, 0xF1.toByte, 0xF1.toByte, 0xF1.toByte, 0xF1.toByte, + 0xF2.toByte, 0xC1.toByte, 0xC1.toByte, 0xC1.toByte, 0xC1.toByte, 0xC1.toByte, + 0xF3.toByte, 0xF3.toByte, 0xF3.toByte, 0xF3.toByte, 0xF3.toByte, 0xF3.toByte, + 0xF4.toByte, 0xC2.toByte, 0xC2.toByte, 0xC2.toByte, 0xC2.toByte, 0xC2.toByte + ) + + assertArraysEqual(bytes, expected) + } + } + + "write using the third alternative of a three-way REDEFINES chain" in { + withTempDirectory("cobol_writer_redefines") { tempDir => + val copybookContentsWithThreeWayRedefines = + """ 01 RECORD. + 05 A PIC X(1). + 05 B PIC 9(5). + 05 B1 PIC X(5) REDEFINES B. + 05 B2 PIC 9(3)V99 REDEFINES B. + """ + + val df = List(("X", new java.math.BigDecimal("123.45"))).toDF("A", "B2") + + val path = new Path(tempDir, "writer_redefines_third") + + df.coalesce(1) + .write + .format("cobol") + .mode(SaveMode.Overwrite) + .option("copybook_contents", copybookContentsWithThreeWayRedefines) + .save(path.toString) + + val bytes = readPartFileBytes(path) + + // Expected EBCDIC data: A='X', B2=123.45 (DISPLAY digits, implied decimal point) + val expected = Array[Byte]( + 0xE7.toByte, + 0xF1.toByte, 0xF2.toByte, 0xF3.toByte, 0xF4.toByte, 0xF5.toByte + ) + + assertArraysEqual(bytes, expected) + } + } + + "fail fast when two non-adjacent alternatives of a three-way REDEFINES chain are both populated" in { + withTempDirectory("cobol_writer_redefines") { tempDir => + val copybookContentsWithThreeWayRedefines = + """ 01 RECORD. + 05 A PIC X(1). + 05 B PIC 9(5). + 05 B1 PIC X(5) REDEFINES B. + 05 B2 PIC 9(3)V99 REDEFINES B. + """ + + val df = List(("X", 12345, new java.math.BigDecimal("123.45"))).toDF("A", "B", "B2") + + val path = new Path(tempDir, "writer_redefines_third_conflict") + + val thrown = intercept[Throwable] { + df.coalesce(1) + .write + .format("cobol") + .mode(SaveMode.Overwrite) + .option("copybook_contents", copybookContentsWithThreeWayRedefines) + .option("write_strict_redefines", "true") + .save(path.toString) + } + + val messages = causeChainMessages(thrown) + assert(messages.exists(m => m.contains("B") && m.contains("B2")), + s"Expected an error mentioning both conflicting REDEFINES fields 'B' and 'B2', but got: ${messages.mkString(" | ")}") + } + } + + "write the first alternative of a three-way REDEFINES chain when multiple are populated and strict is disabled" in { + withTempDirectory("cobol_writer_redefines") { tempDir => + val copybookContentsWithThreeWayRedefines = + """ 01 RECORD. + 05 A PIC X(1). + 05 B PIC 9(5). + 05 B1 PIC X(5) REDEFINES B. + 05 B2 PIC 9(3)V99 REDEFINES B. + """ + + val df = List(("X", 12345, new java.math.BigDecimal("123.45"))).toDF("A", "B", "B2") + + val path = new Path(tempDir, "writer_redefines_three_first_wins") + + df.coalesce(1) + .write + .format("cobol") + .mode(SaveMode.Overwrite) + .option("copybook_contents", copybookContentsWithThreeWayRedefines) + .save(path.toString) + + val bytes = readPartFileBytes(path) + + // A='X' (0xE7), then B=12345 (0xF1..0xF5). B2 is ignored. + val expected = Array[Byte]( + 0xE7.toByte, + 0xF1.toByte, 0xF2.toByte, 0xF3.toByte, 0xF4.toByte, 0xF5.toByte + ) + + assertArraysEqual(bytes, expected) + } + } + + "write zero bytes when none of the REDEFINES alternatives are present and strict schema is disabled" in { + withTempDirectory("cobol_writer_redefines") { tempDir => + val df = List(Tuple1("X")).toDF("A") + + val path = new Path(tempDir, "writer_redefines_none_present") + + df.coalesce(1) + .write + .format("cobol") + .mode(SaveMode.Overwrite) + .option("copybook_contents", copybookContentsWithRedefines) + .option("strict_schema", "false") + .save(path.toString) + + val bytes = readPartFileBytes(path) + + // Expected EBCDIC data: A='X', shared bytes left as zeroes + val expected = Array[Byte]( + 0xE7.toByte, + 0x00, 0x00, 0x00, 0x00, 0x00 + ) + + assertArraysEqual(bytes, expected) + } + } + + "fail with a clear message when none of the REDEFINES alternatives are present and strict schema is enabled" in { + withTempDirectory("cobol_writer_redefines") { tempDir => + val df = List(Tuple1("X")).toDF("A") + + val path = new Path(tempDir, "writer_redefines_none_present_strict") + + val ex = intercept[IllegalArgumentException] { + df.coalesce(1) + .write + .format("cobol") + .mode(SaveMode.Overwrite) + .option("copybook_contents", copybookContentsWithRedefines) + .save(path.toString) + } + + assert(ex.getMessage.contains("B")) + assert(ex.getMessage.contains("B1")) + } + } + + "write the full width of the REDEFINES cluster (max alternative size) and zero-pad unused trailing bytes " + + "when alternatives have different sizes" in { + withTempDirectory("cobol_writer_redefines") { tempDir => + val copybookContentsDifferentSizes = + """ 01 RECORD. + 05 A PIC X(1). + 05 B PIC 9(5). + 05 B1 PIC X(35) REDEFINES B. + """ + + val df = List(("X", 12345, Option.empty[String])).toDF("A", "B", "B1") + + val path = new Path(tempDir, "writer_redefines_diff_sizes") + + df.coalesce(1) + .write + .format("cobol") + .mode(SaveMode.Overwrite) + .option("copybook_contents", copybookContentsDifferentSizes) + .save(path.toString) + + val bytes = readPartFileBytes(path) + + // The REDEFINES cluster always reserves the widest alternative's size (35 bytes for B1), + // even when the narrower alternative (B, 5 bytes) is the one populated. + // Expected EBCDIC data: A='X', B=12345 (DISPLAY digits), followed by 30 zero-bytes + // (the unused tail of B1's 35-byte region, left as binary zeroes rather than spaces). + val expected = Array[Byte](0xE7.toByte, 0xF1.toByte, 0xF2.toByte, 0xF3.toByte, 0xF4.toByte, 0xF5.toByte) ++ + Array.fill[Byte](30)(0x00) + + assert(bytes.length == 36, s"Expected a 36-byte record (1 + max(5, 35)), got ${bytes.length}") + assertArraysEqual(bytes, expected) + } + } + + "write REDEFINES groups that contain nested sub-fields, selecting the active alternative via a " + + "record-type discriminator column" in { + withTempDirectory("cobol_writer_redefines") { tempDir => + // REC-TYPE is a plain data column: it carries no special meaning to the writer (the active + // alternative is still chosen purely by which group is non-null on a given row), but it follows + // the common COBOL convention of tagging each record with a type code so that a reader can tell + // which REDEFINES alternative was used to write a row without inspecting the group contents. + val copybookContentsWithNestedRedefines = + """ 01 RECORD. + 05 A PIC X(1). + 05 REC-TYPE PIC X(1). + 05 GRP-B. + 10 B-NUM PIC 9(5). + 10 B-NAME PIC X(5). + 05 GRP-B1 REDEFINES GRP-B. + 10 B1-CODE PIC X(3). + 10 B1-AMT PIC 9(7). + """ + + val exampleJsons = Seq( + """{"A":"1","REC_TYPE":"B","GRP_B":{"B_NUM":12345,"B_NAME":"HELLO"}}""", + """{"A":"2","REC_TYPE":"1","GRP_B1":{"B1_CODE":"XYZ","B1_AMT":9876543}}""" + ) + + val df = spark.read.json(exampleJsons.toDS()) + .select("A", "REC_TYPE", "GRP_B", "GRP_B1") + + val path = new Path(tempDir, "writer_redefines_nested_groups") + + df.coalesce(1) + .orderBy("A") + .write + .format("cobol") + .mode(SaveMode.Overwrite) + .option("copybook_contents", copybookContentsWithNestedRedefines) + .save(path.toString) + + val bytes = readPartFileBytes(path) + + // Row 1: A='1', REC-TYPE='B', GRP-B populated (B-NUM=12345, B-NAME="HELLO"), GRP-B1 absent. + val row1 = Array[Byte]( + 0xF1.toByte, 0xC2.toByte, + 0xF1.toByte, 0xF2.toByte, 0xF3.toByte, 0xF4.toByte, 0xF5.toByte, + 0xC8.toByte, 0xC5.toByte, 0xD3.toByte, 0xD3.toByte, 0xD6.toByte + ) + // Row 2: A='2', REC-TYPE='1', GRP-B1 populated (B1-CODE="XYZ", B1-AMT=9876543), GRP-B absent. + val row2 = Array[Byte]( + 0xF2.toByte, 0xF1.toByte, + 0xE7.toByte, 0xE8.toByte, 0xE9.toByte, + 0xF9.toByte, 0xF8.toByte, 0xF7.toByte, 0xF6.toByte, 0xF5.toByte, 0xF4.toByte, 0xF3.toByte + ) + + assertArraysEqual(bytes, row1 ++ row2) + } + } + } + } + + def readPartFileBytes(path: Path): Array[Byte] = { + val fs = path.getFileSystem(spark.sparkContext.hadoopConfiguration) + + assert(fs.exists(path), "Output directory should exist") + val files = fs.listStatus(path) + .filter(_.getPath.getName.startsWith("part-")) + assert(files.nonEmpty, "Output directory should contain part files") + + val partFile = files.head.getPath + val data = fs.open(partFile) + val bytes = new Array[Byte](files.head.getLen.toInt) + data.readFully(bytes) + data.close() + bytes + } + + def causeChainMessages(t: Throwable): List[String] = { + @tailrec + def loop(current: Throwable, acc: List[String], seen: Set[Throwable]): List[String] = { + if (current == null || seen.contains(current)) { + acc + } else { + val message = Option(current.getMessage).getOrElse("") + loop(current.getCause, acc :+ message, seen + current) + } + } + loop(t, Nil, Set.empty) }