Retain metadata when mapping root level fields.

AbsaOSS · Jun 3, 2024 · b068a90 · b068a90
1 parent 846d76f
commit b068a90
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/utils/SparkUtils.scala b/spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/utils/SparkUtils.scala
@@ -184,11 +184,11 @@ object SparkUtils extends Logging {
       field.dataType match {
         case st: StructType =>
           val columns = st.fields.map(f => mapField(column.getField(field.name), f))
-          struct(columns: _*).as(field.name)
+          struct(columns: _*).as(field.name, field.metadata)
         case ar: ArrayType =>
-          mapArray(ar, column, field.name).as(field.name)
+          mapArray(ar, column, field.name).as(field.name, field.metadata)
         case _ =>
-          f(field, column).as(field.name)
+          f(field, column).as(field.name, field.metadata)
       }
     }
 

diff --git a/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/utils/SparkUtilsSuite.scala b/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/utils/SparkUtilsSuite.scala
@@ -451,7 +451,6 @@ class SparkUtilsSuite extends AnyFunSuite with SparkTestBase with BinaryFileFixt
         |""".stripMargin
 
     if (!Properties.versionNumberString.startsWith("2.11.")) {
-      println(Properties.versionNumberString)
       withTempTextFile("flatten", "test", StandardCharsets.UTF_8, "") { filePath =>
         val df = spark.read
           .format("cobol")
@@ -461,11 +460,13 @@ class SparkUtilsSuite extends AnyFunSuite with SparkTestBase with BinaryFileFixt
           .option("metadata", "extended")
           .load(filePath)
 
-
         // This method only works with Scala 2.12+ and Spark 3.0+
         val actualDf = SparkUtils.covertIntegralToDecimal(df)
         val actualSchema = actualDf.schema.treeString
 
+        assert(actualDf.schema.fields.head.metadata.json.nonEmpty)
+        assert(actualDf.schema.fields(1).metadata.json.nonEmpty)
+
         compareText(actualSchema, expectedSchema)
       }
     }