Merge pull request #836 from capitalone-contributions/multisegment-header-trailer

Multisegment Header Trailer Removal

Author: yruslan

.gitignore

@@ -49,3 +49,4 @@ output
 #  syntax: regexp
 #  ^\.pc/
 
+.vscode/

README.md

@@ -1642,7 +1642,8 @@ The output looks like this:
 | .option("enable_index_cache", "true")                                                 | When true (default), calculated indexes are cached in memory for later use. This improves performance of processing when same files are processed more than once.                                                                                                                                                                                                                  |
 | .option("input_split_records", 50000)                                                 | Specifies how many records will be allocated to each split/partition. It will be processed by Spark tasks. (The default is not set and the split will happen according to size, see the next option)                                                                                                                                                                               |
 | .option("input_split_size_mb", 100)                                                   | Specify how many megabytes to allocate to each partition/split. (The default is 100 MB)                                                                                                                                                                                                                                                                                            |
-
+| .option("record_header_name", "HEADER")                                               | Assuming a copybook definition represents a header, offsets the file read start by the number of bytes in that header and excludes the record definition from the output schema.                                                                                                                                                                                                                                                                                           |
+| .option("record_trailer_name", "TRAILER")                                             | Assuming a copybook definition represents a trailer, offsets the file read end by the number of bytes in that trailer and excludes the record definition from the output schema.                                                                                                                                                                                                                                                                                            |
 ##### Helper fields generation options    
 
 | Option (usage example)                     | Description                                                                                                                                                                         |

cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/extractors/record/RecordExtractors.scala

@@ -71,7 +71,8 @@ object RecordExtractors {
                                   activeSegmentRedefine: String = "",
                                   generateInputFileField: Boolean = false,
                                   inputFileName: String = "",
-                                  handler: RecordHandler[T]
+                                  handler: RecordHandler[T],
+                                  recordsToExclude: Set[String] = Set.empty
   ): Seq[Any] = {
     val dependFields = scala.collection.mutable.HashMap.empty[String, Either[Int, String]]
     val corruptFields = new ArrayBuffer[CorruptField]
@@ -202,7 +203,7 @@ object RecordExtractors {
 
     val records: ListBuffer[T] = ListBuffer.empty[T]
 
-    for (record <- rootRecords) yield {
+    for (record <- rootRecords if !recordsToExclude.contains(record.name.toUpperCase)) yield {
       val (size, values) = getGroupValues(nextOffset, record.asInstanceOf[Group])
       if (!record.isRedefined) {
         nextOffset += size
@@ -259,7 +260,8 @@ object RecordExtractors {
       recordId: Long = 0,
       generateInputFileField: Boolean = false,
       inputFileName: String = "",
-      handler: RecordHandler[T]
+      handler: RecordHandler[T],
+      recordsToExclude: Set[String] = Set.empty
   ): Seq[Any] = {
     val isAstFlat = ast.children.exists(_.isInstanceOf[Primitive])
 
@@ -422,7 +424,7 @@ object RecordExtractors {
 
     val records: ListBuffer[T] = ListBuffer.empty[T]
 
-    rootRecords.collect { case grp: Group if grp.parentSegment.isEmpty =>
+    rootRecords.collect { case grp: Group if grp.parentSegment.isEmpty && !recordsToExclude.contains(grp.name.toUpperCase) =>
       val (size, values) = getGroupValues(nextOffset, grp, segmentsData(0)._2, 0, segmentsData(0)._1 :: Nil)
       nextOffset += size
       records += values

cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/FixedLenNestedRowIterator.scala

@@ -96,7 +96,8 @@ class FixedLenNestedRowIterator[T: ClassTag](
       generateCorruptFields = generateCorruptFields,
       generateCorruptFieldsAsHex = generateCorruptFieldsAsHex,
       activeSegmentRedefine = activeSegmentRedefine,
-      handler = handler
+      handler = handler,
+      recordsToExclude = readerProperties.recordsToExclude
     )
 
     // Advance byte index to the next record

cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VarLenHierarchicalIterator.scala

@@ -149,7 +149,8 @@ final class VarLenHierarchicalIterator[T: ClassTag](cobolSchema: Copybook,
       recordIndex,
       generateInputFileName,
       dataStream.inputFileName,
-      handler
+      handler,
+      readerProperties.recordsToExclude
     )
   }
 

cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VarLenNestedIterator.scala

@@ -109,7 +109,8 @@ final class VarLenNestedIterator[T: ClassTag](cobolSchema: Copybook,
                 activeSegmentRedefine = segmentRedefine,
                 generateInputFileName,
                 dataStream.inputFileName,
-                handler
+                handler,
+                readerProperties.recordsToExclude
               ))
 
               recordFetched = true

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

@@ -111,5 +111,7 @@ case class CobolParameters(
                             debugLayoutPositions:    Boolean,
                             enableSelfChecks:        Boolean,
                             metadataPolicy:          MetadataPolicy,
+                            recordHeaderName:        Option[String] = None,
+                            recordTrailerName:       Option[String] = None,
                             options:                 Map[String, String]
                           )

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

@@ -60,6 +60,8 @@ object CobolParametersParser extends Logging {
   val PARAM_RECORD_END_OFFSET         = "record_end_offset"
   val PARAM_FILE_START_OFFSET         = "file_start_offset"
   val PARAM_FILE_END_OFFSET           = "file_end_offset"
+  val PARAM_RECORD_HEADER_NAME        = "record_header_name"
+  val PARAM_RECORD_TRAILER_NAME       = "record_trailer_name"
   val PARAM_IS_XCOM                   = "is_xcom"
   val PARAM_IS_TEXT                   = "is_text"
 
@@ -310,6 +312,8 @@ object CobolParametersParser extends Logging {
       params.getOrElse(PARAM_DEBUG_LAYOUT_POSITIONS, "false").toBoolean,
       params.getOrElse(PARAM_ENABLE_SELF_CHECKS, "false").toBoolean,
       MetadataPolicy(params.getOrElse(PARAM_METADATA, "basic")),
+      params.get(PARAM_RECORD_HEADER_NAME).map(_.trim).filter(_.nonEmpty),
+      params.get(PARAM_RECORD_TRAILER_NAME).map(_.trim).filter(_.nonEmpty),
       params.getMap
       )
     validateSparkCobolOptions(params, recordFormat, validateRedundantOptions)
@@ -412,6 +416,9 @@ object CobolParametersParser extends Logging {
       CorruptFieldsPolicy.Disabled
     }
 
+    val recordsToExclude = (parameters.recordHeaderName.map(n => CopybookParser.transformIdentifier(n).toUpperCase).toSet ++
+      parameters.recordTrailerName.map(n => CopybookParser.transformIdentifier(n).toUpperCase).toSet)
+
     ReaderParameters(
       recordFormat = parameters.recordFormat,
       isEbcdic = parameters.isEbcdic,
@@ -470,6 +477,7 @@ object CobolParametersParser extends Logging {
       varLenParams.reAdditionalInfo,
       varLenParams.inputFileNameColumn,
       parameters.metadataPolicy,
+      recordsToExclude,
       parameters.options
       )
   }
@@ -982,6 +990,22 @@ object CobolParametersParser extends Logging {
       params.contains(PARAM_ENABLE_INDEX_CACHE) && params(PARAM_ENABLE_INDEX_CACHE).toBoolean)
       throw new IllegalArgumentException(s"When '$PARAM_ENABLE_INDEXES' = false, '$PARAM_ENABLE_INDEX_CACHE' cannot be true.")
 
+    if (params.contains(PARAM_RECORD_HEADER_NAME) && params.contains(PARAM_FILE_START_OFFSET)) {
+      throw new IllegalArgumentException(s"Options '$PARAM_RECORD_HEADER_NAME' and '$PARAM_FILE_START_OFFSET' cannot be used together.")
+    }
+
+    if (params.contains(PARAM_RECORD_TRAILER_NAME) && params.contains(PARAM_FILE_END_OFFSET)) {
+      throw new IllegalArgumentException(s"Options '$PARAM_RECORD_TRAILER_NAME' and '$PARAM_FILE_END_OFFSET' cannot be used together.")
+    }
+
+    if (params.contains(PARAM_RECORD_HEADER_NAME) && params.contains(PARAM_RECORD_TRAILER_NAME)) {
+      val headerName = params(PARAM_RECORD_HEADER_NAME).trim
+      val trailerName = params(PARAM_RECORD_TRAILER_NAME).trim
+      if (headerName.equalsIgnoreCase(trailerName)) {
+        throw new IllegalArgumentException(s"Options '$PARAM_RECORD_HEADER_NAME' and '$PARAM_RECORD_TRAILER_NAME' cannot refer to the same record ('$headerName').")
+      }
+    }
+
     if (validateRedundantOptions && unusedKeys.nonEmpty) {
       val unusedKeyStr = unusedKeys.mkString(",")
       val msg = s"Redundant or unrecognized option(s) to 'spark-cobol': $unusedKeyStr."

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

@@ -139,5 +139,6 @@ case class ReaderParameters(
                              reAdditionalInfo:        String = "",
                              inputFileNameColumn:     String = "",
                              metadataPolicy:          MetadataPolicy = MetadataPolicy.Basic,
+                             recordsToExclude:       Set[String] = Set.empty,
                              options:                 Map[String, String] = Map.empty
                            )

cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/schema/CobolSchema.scala

@@ -55,13 +55,24 @@ class CobolSchema(val copybook: Copybook,
                   val corruptSchemaPolicy: CorruptFieldsPolicy,
                   val generateSegIdFieldsCnt: Int = 0,
                   val segmentIdProvidedPrefix: String = "",
-                  val metadataPolicy: MetadataPolicy = MetadataPolicy.Basic) extends Serializable {
+                  val metadataPolicy: MetadataPolicy = MetadataPolicy.Basic,
+                  val recordsToExclude: Set[String] = Set.empty) extends Serializable {
 
   val segmentIdPrefix: String = if (segmentIdProvidedPrefix.isEmpty) getDefaultSegmentIdPrefix else segmentIdProvidedPrefix
 
   def getCobolSchema: Copybook = copybook
 
-  lazy val getRecordSize: Int = copybook.getRecordSize
+  lazy val getRecordSize: Int = {
+    if (recordsToExclude.isEmpty) {
+      copybook.getRecordSize
+    } else {
+      val excludedSize = copybook.ast.children
+        .filter(r => recordsToExclude.contains(r.name.toUpperCase))
+        .map(_.binaryProperties.actualSize)
+        .sum
+      copybook.getRecordSize - excludedSize
+    }
+  }
 
   def isRecordFixedSize: Boolean = copybook.isRecordFixedSize
 
@@ -147,7 +158,8 @@ object CobolSchema {
       readerParameters.corruptFieldsPolicy,
       segIdFieldCount,
       segmentIdPrefix,
-      readerParameters.metadataPolicy
+      readerParameters.metadataPolicy,
+      readerParameters.recordsToExclude
     )
   }