#863 Add support for custom AST transformers via the ast_transformers option.

README.md

@@ -1635,6 +1635,49 @@ The output looks like this:
 +---+-------+----------+----------+----------+
 ```
 
+### Custom AST transformers
+When parsed copybooks need to be transformed in a regular way before processing data, a custom AST transformer can be
+used. This allows modifying data types on the fly, or adding/removing fields when necessary.
+
+In order to use a custom AST transformer, you need to implement the `AstTransformer` trait and provide your 
+transformation logic in the `transform` method. Then, you can specify your custom transformer class in the `ast_transformers`
+option when reading data.
+
+Example transformer:
+
+```scala
+package com.example
+
+import za.co.absa.cobrix.cobol.parser.asttransform.AstTransformer
+import za.co.absa.cobrix.cobol.reader.parameters.ReaderParameters
+import za.co.absa.cobrix.cobol.parser.CopybookParser.CopybookAST
+
+class MyTransformer1(readerParameters: ReaderParameters) extends AstTransformer {
+  override def transform(ast: CopybookAST): CopybookAST = {
+    // implementation goes here
+  }
+}
+```
+
+Example options that uses the transformer:
+```scala
+spark.read
+  .format("cobol")
+  .option("copybook", someCopybook)
+  .option("ast_transformers", "com.example.MyTransformer1")
+  .load("/some/path")
+```
+
+You can chain multiple transformers by providing transformer classes as a comma-separated list.
+
+**Notes:**
+
+1. Transformers are applied before built-in transformers, so the AST might not be complete at that stage.
+2. Since the decoder and optionally an encoder are set at the AST building stage, if you want to insert new fields,
+   you need to define them as well. Use `DecoderSelector.getDecoder()` and [optionally] `EncoderSelector.getEncoder()`.
+3. Custom transformers run before binary properties are calculated. So inserting or removing fields changes the copybook
+   layout.
+
 ## Summary of all available options
 
 ##### File reading options
@@ -1688,6 +1731,7 @@ The output looks like this:
 | .option("record_limit", "1000")                     | Caps decoded output to N rows globally across all input paths/files (zero returns no rows). This is a source-level global cap: `df.limit(N)` does not push down into Cobrix. While enabled Cobrix uses a single scan task and can avoid opening later input partitions/files once satisfied. Indexed reads may still build indexes first. |
 | .option("with_input_file_name_col", "file_name")    | Generates a column containing input file name for each record (Similar to Spark SQL `input_file_name()` function). The column name is specified by the value of the option. This option only works for variable record length files. For fixed record length and ASCII files use `input_file_name()`.                                     |
 | .option("metadata", "basic")                        | Specifies wat kind of metadata to include in the Spark schema: `false`, `basic`(default), or `extended` (PIC, usage, etc).                                                                                                                                                                                                                |
+| .option("ast_transformers", "...")                  | Specifies a comma-separated list of custom AST transformers to be used to process AST of copybooks just after they are parsed.                                                                                                                                                                                                            |
 | .option("debug", "hex")                             | If specified, each primitive field will be accompanied by a debug field containing raw bytes from the source file. Possible values: `none` (default), `hex`, `binary`, `string` (ASCII only). The legacy value `true` is supported and will generate debug fields in HEX.                                                                 |
 
 For example, to cap a Cobrix source read at 1,000 decoded rows:
@@ -1815,7 +1859,7 @@ val df = spark
 
 `common_extended`, `cp037_extended` are code pages supporting non-printable characters that converts to ASCII codes below 32.
 
-## EBCDIC Processor (experimental)
+## EBCDIC Processor
 The EBCDIC processor allows processing files by replacing value of fields without changing the underlying format (`CobolProcessingStrategy.InPlace`)
 or with conversion of the input format to variable-record-length format with big-endian RDWs (`CobolProcessingStrategy.ToVariableLength`).
 
@@ -1871,7 +1915,7 @@ val count = CobolProcessor.builder
 ```
 
 
-## EBCDIC Spark Processor (experimental)
+## EBCDIC Spark Processor
 This allows in-place processing of data retaining original format in parallel uring RDDs under the hood.
 
 Here is an example usage:
@@ -1898,7 +1942,7 @@ SparkCobolProcessor.builder
   .save(outputPath)
 ```
 
-## EBCDIC Spark raw record RDD generator (experimental)
+## EBCDIC Spark raw record RDD generator
 You can process raw records of a mainframe file as an `RDD[Array[Byte]]`. This can be useful for custom processing without converting
 to Spark data types. You can still access fields via parsed copybooks.
 
@@ -1927,7 +1971,7 @@ val segmentRdds RDD[String] = recordsRdd.flatMap { record =>
 segmentRdds.distinct.collect.sorted.foreach(println)
 ```
 
-## EBCDIC Writer (experimental)
+## EBCDIC Writer
 
 Cobrix's EBCDIC writer is an experimental feature that allows writing Spark DataFrames as EBCDIC mainframe files.
 

cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/parser/CopybookParser.scala

@@ -127,6 +127,7 @@ object CopybookParser extends Logging {
     * @param nonTerminals          A list of non-terminals that should be extracted as strings.
     * @param redefineRuleExpressions A map of REDEFINE field names to expressions that determine which redefine alternative to use when parsing records.
     * @param debugFieldsPolicy     Specifies if debugging fields need to be added and what should they contain (false, hex, raw).
+    * @param customAstTransformers A list of custom AST transformers to be applied during parsing.
     * @return Seq[Group] where a group is a record inside the copybook.
     */
   def parse(copyBookContents: String,
@@ -151,7 +152,8 @@ object CopybookParser extends Logging {
             occursHandlers: Map[String, Map[String, Int]] = Map(),
             redefineRuleExpressions: Map[String, ExpressionEvaluator] = Map.empty,
             debugFieldsPolicy: DebugFieldsPolicy = DebugFieldsPolicy.NoDebug,
-            fieldCodePageMap: Map[String, String] = Map.empty[String, String]): Copybook = {
+            fieldCodePageMap: Map[String, String] = Map.empty[String, String],
+            customAstTransformers: Seq[AstTransformer] = Seq.empty): Copybook = {
     parseTree(dataEncoding,
       copyBookContents,
       dropGroupFillers,
@@ -174,7 +176,8 @@ object CopybookParser extends Logging {
       occursHandlers,
       redefineRuleExpressions,
       debugFieldsPolicy,
-      fieldCodePageMap)
+      fieldCodePageMap,
+      customAstTransformers)
   }
 
   /**
@@ -198,6 +201,7 @@ object CopybookParser extends Logging {
     * @param nonTerminals          A list of non-terminals that should be extracted as strings
     * @param redefineRuleExpressions A map of REDEFINE field names to expressions that determine which redefine alternative to use when parsing records.
     * @param debugFieldsPolicy     Specifies if debugging fields need to be added and what should they contain (false, hex, raw).
+    * @param customAstTransformers A list of custom AST transformers to be applied during parsing.
     * @return Seq[Group] where a group is a record inside the copybook
     */
   def parseTree(copyBookContents: String,
@@ -221,7 +225,8 @@ object CopybookParser extends Logging {
                 occursHandlers: Map[String, Map[String, Int]] = Map(),
                 redefineRuleExpressions: Map[String, ExpressionEvaluator] = Map.empty,
                 debugFieldsPolicy: DebugFieldsPolicy = DebugFieldsPolicy.NoDebug,
-                fieldCodePageMap: Map[String, String] = Map.empty[String, String]): Copybook = {
+                fieldCodePageMap: Map[String, String] = Map.empty[String, String],
+                customAstTransformers: Seq[AstTransformer] = Seq.empty): Copybook = {
     parseTree(EBCDIC,
       copyBookContents,
       dropGroupFillers,
@@ -244,7 +249,8 @@ object CopybookParser extends Logging {
       occursHandlers,
       redefineRuleExpressions,
       debugFieldsPolicy,
-      fieldCodePageMap)
+      fieldCodePageMap,
+      customAstTransformers)
   }
 
   /**
@@ -268,6 +274,7 @@ object CopybookParser extends Logging {
     * @param nonTerminals          A list of non-terminals that should be extracted as strings
     * @param redefineRuleExpressions A map of REDEFINE field names to expressions that determine which redefine alternative to use when parsing records.
     * @param debugFieldsPolicy     Specifies if debugging fields need to be added and what should they contain (false, hex, raw).
+    * @param customAstTransformers A list of custom AST transformers to be applied during parsing.
     * @return Seq[Group] where a group is a record inside the copybook
     */
   @throws(classOf[SyntaxErrorException])
@@ -293,7 +300,8 @@ object CopybookParser extends Logging {
                 occursHandlers: Map[String, Map[String, Int]],
                 redefineRuleExpressions: Map[String, ExpressionEvaluator],
                 debugFieldsPolicy: DebugFieldsPolicy,
-                fieldCodePageMap: Map[String, String]): Copybook = {
+                fieldCodePageMap: Map[String, String],
+                customAstTransformers: Seq[AstTransformer]): Copybook = {
 
     val schemaANTLR: CopybookAST = ANTLRParser.parse(copyBookContents, enc, stringTrimmingPolicy, isDisplayAlwaysString, commentPolicy, strictSignOverpunch, improvedNullDetection, strictIntegralPrecision, decodeBinaryAsHex, ebcdicCodePage, asciiCharset, isUtf16BigEndian, floatingPointFormat, fieldCodePageMap)
 
@@ -304,7 +312,7 @@ object CopybookParser extends Logging {
     val correctedFieldParentMap = transformIdentifierMap(fieldParentMap)
     validateFieldParentMap(correctedFieldParentMap)
 
-    val transformers = Seq(
+    val transformers = customAstTransformers ++ Seq(
       // Calculate sizes of fields and their positions from the beginning of a record
       BinaryPropertiesAdder(),
       // Adds virtual primitive fields for GROUPs that can be parsed as concatenation of their children.

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

@@ -69,6 +69,7 @@ import za.co.absa.cobrix.cobol.reader.policies.SchemaRetentionPolicy.SchemaReten
   * @param writerParameters        Parameters set only for the writer of files in mainfraame format
   * @param options                 Options passed to 'spark-cobol'.
   * @param recordLimit             Maximum number of decoded output rows to return. Zero returns no rows; if not specified, all rows are returned.
+  * @param astTransformerClasses   A comma-separated list of fully qualified classes extending [[za.co.absa.cobrix.cobol.parser.asttransform.AstTransformer]]
   */
 case class CobolParameters(
                             copybookPath:            Option[String],
@@ -118,6 +119,7 @@ case class CobolParameters(
                             fileHeaderField:         Option[String] = None,
                             fileTrailerField:        Option[String] = None,
                             recordLimit:             Option[Int] = None,
+                            astTransformerClasses:   Seq[String],
                             writerParameters:        Option[WriterParameters],
                             options:                 Map[String, String]
                           )

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

@@ -80,6 +80,7 @@ object CobolParametersParser extends Logging {
   val PARAM_STRICT_INTEGRAL_PRECISION = "strict_integral_precision"
   val PARAM_DISPLAY_PIC_ALWAYS_STRING = "display_pic_always_string"
   val PARAM_STRICT_SCHEMA             = "strict_schema"
+  val PARAM_AST_TRANSFORMERS          = "ast_transformers"
 
   val PARAM_GROUP_NOT_TERMINALS       = "non_terminals"
   val PARAM_OCCURS_MAPPINGS           = "occurs_mappings"
@@ -309,6 +310,9 @@ object CobolParametersParser extends Logging {
       None
     }
 
+    val astTransformerClasses = params.get(PARAM_AST_TRANSFORMERS)
+      .map(_.split(',').map(_.trim).filter(_.nonEmpty).toSeq).getOrElse(Seq.empty)
+
     val cobolParameters = CobolParameters(
       copybookPathOpt,
       copybookPaths,
@@ -357,6 +361,7 @@ object CobolParametersParser extends Logging {
       params.get(PARAM_RECORD_HEADER_NAME).orElse(params.get(PARAM_RECORD_HEADER_NAME2)) .map(_.trim).filter(_.nonEmpty),
       params.get(PARAM_RECORD_TRAILER_NAME).orElse(params.get(PARAM_RECORD_TRAILER_NAME2)) .map(_.trim).filter(_.nonEmpty),
       recordLimit,
+      astTransformerClasses,
       writerParameters,
       params.getMap
       )
@@ -540,6 +545,7 @@ object CobolParametersParser extends Logging {
       parameters.metadataPolicy,
       recordsToExclude,
       parameters.recordLimit,
+      parameters.astTransformerClasses,
       parameters.writerParameters,
       parameters.options
       )

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

@@ -82,6 +82,7 @@ import za.co.absa.cobrix.cobol.reader.policies.SchemaRetentionPolicy.SchemaReten
   * @param metadataPolicy          Specifies the policy of metadat fields to be added to the Spark schema
   * @param recordsToExclude        A set of fields to exclude from decoding in normal (not header or footer) records.
   * @param recordLimit             Maximum number of decoded output rows to return. Zero returns no rows; if not specified, all rows are returned.
+  * @param astTransformerClasses   A comma-separated list of fully qualified classes extending [[za.co.absa.cobrix.cobol.parser.asttransform.AstTransformer]]
   * @param writerParameters        If specified, contains parameters for the writer.
   * @param options                 Options passed to spark-cobol
   */
@@ -146,6 +147,7 @@ case class ReaderParameters(
                              metadataPolicy:          MetadataPolicy = MetadataPolicy.Basic,
                              recordsToExclude:        Set[String] = Set.empty,
                              recordLimit:             Option[Int] = None,
+                             astTransformerClasses:   Seq[String] = Seq.empty,
                              writerParameters:        Option[WriterParameters] = None,
                              options:                 Map[String, String] = Map.empty
                            )

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

@@ -22,6 +22,7 @@ import za.co.absa.cobrix.cobol.parser.policies.MetadataPolicy
 import za.co.absa.cobrix.cobol.parser.{Copybook, CopybookParser}
 import za.co.absa.cobrix.cobol.reader.parameters.{CorruptFieldsPolicy, ReaderParameters}
 import za.co.absa.cobrix.cobol.reader.policies.SchemaRetentionPolicy.SchemaRetentionPolicy
+import za.co.absa.cobrix.cobol.utils.AstTransformerUtils
 
 import java.nio.charset.{Charset, StandardCharsets}
 import java.time.ZonedDateTime
@@ -98,6 +99,8 @@ object CobolSchema {
       case None               => StandardCharsets.UTF_8
     }
 
+    val customAstTransformers = readerParameters.astTransformerClasses.map(className => AstTransformerUtils.loadAstTransformer(className, readerParameters))
+
     val schema = if (copyBookContents.size == 1)
       CopybookParser.parseTree(encoding,
         copyBookContents.head,
@@ -121,7 +124,8 @@ object CobolSchema {
         readerParameters.occursMappings,
         readerParameters.redefineRuleExpressions,
         readerParameters.debugFieldsPolicy,
-        readerParameters.fieldCodePage)
+        readerParameters.fieldCodePage,
+        customAstTransformers)
     else
       Copybook.merge(copyBookContents.map(cpb =>
         CopybookParser.parseTree(encoding,
@@ -146,7 +150,8 @@ object CobolSchema {
           readerParameters.occursMappings,
           readerParameters.redefineRuleExpressions,
           readerParameters.debugFieldsPolicy,
-          readerParameters.fieldCodePage)
+          readerParameters.fieldCodePage,
+          customAstTransformers)
       ))
     val segIdFieldCount = readerParameters.multisegment.map(p => p.segmentLevelIds.size).getOrElse(0)
     val segmentIdPrefix = readerParameters.multisegment.map(p => p.segmentIdPrefix).getOrElse("")

cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/utils/AstTransformerUtils.scala

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2018 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.cobrix.cobol.utils
+
+import za.co.absa.cobrix.cobol.parser.asttransform.AstTransformer
+import za.co.absa.cobrix.cobol.reader.parameters.ReaderParameters
+
+import java.lang.reflect.Constructor
+import scala.util.Try
+
+object AstTransformerUtils {
+  /**
+    * Loads and instantiates an AstTransformer by its fully qualified class name using reflection.
+    *
+    * The method attempts to find and invoke a constructor that accepts a ReaderParameters argument.
+    * If no such constructor exists, it falls back to using the default no-argument constructor.
+    *
+    * @param fullyQualifiedName the fully qualified class name of the AstTransformer implementation to load
+    * @param readerParameters   the reader parameters to pass to the constructor if a matching constructor is available
+    * @return an instance of the specified AstTransformer class
+    * @throws java.lang.ClassNotFoundException if the class with the given fully qualified name cannot be found
+    * @throws java.lang.ClassCastException     if the instantiated object cannot be cast to AstTransformer
+    * @throws java.lang.NoSuchMethodException  if neither a constructor accepting ReaderParameters nor a default constructor is found
+    */
+  @throws[ClassNotFoundException]
+  @throws[ClassCastException]
+  @throws[NoSuchMethodException]
+  def loadAstTransformer(fullyQualifiedName: String, readerParameters: ReaderParameters): AstTransformer = {
+    // There are 2 types of constructors supported. If there is a one that takes a ReaderParameters object - use it.
+    // Otherwise, try the default constructor.
+    val confCtor = Try[Constructor[_]](Class.forName(fullyQualifiedName).getConstructor(classOf[ReaderParameters]))
+
+    confCtor
+      .map(ctor => ctor.newInstance(readerParameters).asInstanceOf[AstTransformer])
+      .recoverWith {
+        case _: NoSuchMethodException =>
+          val defCtor = Try[Constructor[_]](Class.forName(fullyQualifiedName).getConstructor())
+          defCtor.map(ctor => ctor.newInstance().asInstanceOf[AstTransformer])
+      }.get
+  }
+}