#857 Add source-level record limit

README.md

@@ -1685,10 +1685,22 @@ The output looks like this:
 | .option("generate_record_id", false)                | Generate autoincremental 'File_Id', 'Record_Id' and 'Record_Byte_Length' fields. This is used for processing record order dependent data.                                                                                                                                                             |
 | .option("generate_record_bytes", false)             | Generate 'Record_Bytes', the binary field that contains raw contents of the original unparsed records.                                                                                                                                                                                                |
 | .option("generate_corrupt_fields", false)           | Generate `_corrupt_fields` field that contains values of fields Cobrix was unable to decode.                                                                                                                                                                                                          |
+| .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("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:
+
+```scala
+val df = spark
+  .read
+  .format("cobol")
+  .option("copybook", "/path/to/copybook")
+  .option("record_limit", "1000")
+  .load("/path/to/data")
+```
+
 ##### Fixed length record format options (for record_format = F or FB)
 
 | Option (usage example)            | Description                                                                                                                                                                                                                                                                             |

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

@@ -68,6 +68,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 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.
   */
 case class CobolParameters(
                             copybookPath:            Option[String],
@@ -117,5 +118,6 @@ case class CobolParameters(
                             fileHeaderField:         Option[String] = None,
                             fileTrailerField:        Option[String] = None,
                             writerParameters:        Option[WriterParameters],
-                            options:                 Map[String, String]
+                            options:                 Map[String, String],
+                            recordLimit:             Option[Int] = None
                           )

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

@@ -54,6 +54,7 @@ object CobolParametersParser extends Logging {
   val PARAM_RECORD_LENGTH             = "record_length"
   val PARAM_MINIMUM_RECORD_LENGTH     = "minimum_record_length"
   val PARAM_MAXIMUM_RECORD_LENGTH     = "maximum_record_length"
+  val PARAM_RECORD_LIMIT              = "record_limit"
   val PARAM_IS_RECORD_SEQUENCE        = "is_record_sequence"
   val PARAM_RECORD_LENGTH_FIELD       = "record_length_field"
   val PARAM_RECORD_LENGTH_MAP         = "record_length_map"
@@ -238,12 +239,32 @@ object CobolParametersParser extends Logging {
     policy
   }
 
+  private def getRecordLimit(params: Parameters): Option[Int] = {
+    params.get(PARAM_RECORD_LIMIT).map { value =>
+      val recordLimit = try {
+        value.toInt
+      } catch {
+        case NonFatal(ex) =>
+          throw new IllegalArgumentException(
+            s"Invalid value '$value' for '$PARAM_RECORD_LIMIT' option. It must be a non-negative 32-bit integer.", ex)
+      }
+
+      if (recordLimit < 0) {
+        throw new IllegalArgumentException(
+          s"Invalid value '$value' for '$PARAM_RECORD_LIMIT' option. It must be a non-negative 32-bit integer.")
+      }
+
+      recordLimit
+    }
+  }
+
   def parse(params: Parameters, validateRedundantOptions: Boolean = true, isWriter: Boolean = false): CobolParameters = {
     val schemaRetentionPolicy = getSchemaRetentionPolicy(params)
     val stringTrimmingPolicy = getStringTrimmingPolicy(params)
     val ebcdicCodePageName = params.getOrElse(PARAM_EBCDIC_CODE_PAGE, "common")
     val ebcdicCodePageClass = params.get(PARAM_EBCDIC_CODE_PAGE_CLASS)
     val asciiCharset = params.get(PARAM_ASCII_CHARSET)
+    val recordLimit = getRecordLimit(params)
 
     val recordFormatDefined = getRecordFormat(params)
 
@@ -334,7 +355,8 @@ 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),
       writerParameters,
-      params.getMap
+      params.getMap,
+      recordLimit
       )
     validateSparkCobolOptions(params, recordFormat, validateRedundantOptions)
     cobolParameters

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

@@ -20,6 +20,57 @@ import org.scalatest.wordspec.AnyWordSpec
 
 class CobolParametersParserSuite extends AnyWordSpec {
   "parse" should {
+    "leave record_limit undefined when it is absent" in {
+      val parsedParams = CobolParametersParser.parse(new Parameters(Map.empty[String, String]))
+
+      assert(parsedParams.recordLimit.isEmpty)
+    }
+
+    "parse a zero record_limit" in {
+      val parsedParams = CobolParametersParser.parse(new Parameters(Map("record_limit" -> "0")))
+
+      assert(parsedParams.recordLimit.contains(0))
+    }
+
+    "parse a positive record_limit" in {
+      val parsedParams = CobolParametersParser.parse(new Parameters(Map("record_limit" -> "100")))
+
+      assert(parsedParams.recordLimit.contains(100))
+    }
+
+    "reject a negative record_limit" in {
+      val exception = intercept[IllegalArgumentException] {
+        CobolParametersParser.parse(new Parameters(Map("record_limit" -> "-1")))
+      }
+
+      assert(exception.getMessage.contains("record_limit"))
+    }
+
+    "reject a malformed record_limit" in {
+      val exception = intercept[IllegalArgumentException] {
+        CobolParametersParser.parse(new Parameters(Map("record_limit" -> "invalid")))
+      }
+
+      assert(exception.getMessage.contains("record_limit"))
+    }
+
+    "reject an overflowing record_limit" in {
+      val exception = intercept[IllegalArgumentException] {
+        CobolParametersParser.parse(new Parameters(Map("record_limit" -> "2147483648")))
+      }
+
+      assert(exception.getMessage.contains("record_limit"))
+    }
+
+    "recognize record_limit in pedantic mode" in {
+      val parsedParams = CobolParametersParser.parse(new Parameters(Map(
+        "record_limit" -> "1",
+        "pedantic" -> "true"
+      )))
+
+      assert(parsedParams.recordLimit.contains(1))
+    }
+
     "parse writer parameters" in {
       val params =  new Parameters(Map(
         "write_null_strings_as_spaces" -> "false",

spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/source/CobolRelation.scala

@@ -66,7 +66,8 @@ class CobolRelation(sourceDirs: Seq[String],
                     filesList: Array[FileWithOrder],
                     cobolReader: Reader,
                     localityParams: LocalityParameters,
-                    debugIgnoreFileSize: Boolean)
+                    debugIgnoreFileSize: Boolean,
+                    recordLimit: Option[Int] = None)
                    (@transient val sqlContext: SQLContext)
   extends BaseRelation
     with Serializable
@@ -79,7 +80,7 @@ class CobolRelation(sourceDirs: Seq[String],
   }
 
   override def buildScan(): RDD[Row] = {
-    cobolReader match {
+    val records = cobolReader match {
       case blockReader: FixedLenTextReader =>
         CobolScanners.buildScanForTextFiles(blockReader, sourceDirs, parseRecords, sqlContext)
       case blockReader: FixedLenReader =>
@@ -91,6 +92,8 @@ class CobolRelation(sourceDirs: Seq[String],
       case _ =>
         throw new IllegalStateException("Invalid reader object $cobolReader.")
     }
+
+    CobolRelation.applyRecordLimit(records, recordLimit)
   }
 
   private[source] def parseRecords(reader: FixedLenReader, records: RDD[Array[Byte]]): RDD[Row] = {
@@ -104,6 +107,14 @@ class CobolRelation(sourceDirs: Seq[String],
 }
 
 object CobolRelation {
+  private[source] def applyRecordLimit(records: RDD[Row], recordLimit: Option[Int]): RDD[Row] = {
+    recordLimit match {
+      case None => records
+      case Some(0) => records.sparkContext.emptyRDD[Row]
+      case Some(limit) => records.coalesce(1, shuffle = false).mapPartitions(_.take(limit))
+    }
+  }
+
   /**
     * Retrieves a list containing the files contained in the directory to be processed attached to numbers which serve
     * as their order.
@@ -128,4 +139,4 @@ object CobolRelation {
         FileWithOrder(fileName, order, isCompressed)
       }
   }
-}
+}

spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/source/DefaultSource.scala

@@ -74,7 +74,8 @@ class DefaultSource
       filesList,
       buildEitherReader(sqlContext.sparkSession, cobolParameters, hasCompressedFiles),
       LocalityParameters.extract(cobolParameters),
-      cobolParameters.debugIgnoreFileSize)(sqlContext)
+      cobolParameters.debugIgnoreFileSize,
+      cobolParameters.recordLimit)(sqlContext)
   }
 
   /** Writer relation */

spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/source/parameters/CobolParametersValidator.scala

@@ -166,6 +166,10 @@ object CobolParametersValidator {
       issues += "Multi-segment options ('segment_field', 'segment_filter', etc) are not supported for writing"
     }
 
+    if (readerParameters.options.contains(PARAM_RECORD_LIMIT)) {
+      issues += s"'$PARAM_RECORD_LIMIT' is supported only for reading"
+    }
+
     if (issues.nonEmpty) {
       throw new IllegalArgumentException(s"Writer validation issues: ${issues.mkString("; ")}")
     }

spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/source/CobolRelationSpec.scala

@@ -69,7 +69,8 @@ class CobolRelationSpec extends SparkCobolTestBase with Serializable {
       filesList,
       testReader,
       localityParams = localityParams,
-      debugIgnoreFileSize = false)(sqlContext)
+      debugIgnoreFileSize = false,
+      recordLimit = None)(sqlContext)
     val cobolData: RDD[Row] = relation.parseRecords(testReader, oneRowRDD)
 
     val cobolDataFrame = sqlContext.createDataFrame(cobolData, sparkSchema)
@@ -96,7 +97,8 @@ class CobolRelationSpec extends SparkCobolTestBase with Serializable {
       filesList,
       testReader,
       localityParams = localityParams,
-      debugIgnoreFileSize = false)(sqlContext)
+      debugIgnoreFileSize = false,
+      recordLimit = None)(sqlContext)
 
     val caught = intercept[Exception] {
       relation.parseRecords(testReader, oneRowRDD).collect()
@@ -114,7 +116,8 @@ class CobolRelationSpec extends SparkCobolTestBase with Serializable {
       filesList,
       testReader,
       localityParams = localityParams,
-      debugIgnoreFileSize = false)(sqlContext)
+      debugIgnoreFileSize = false,
+      recordLimit = None)(sqlContext)
 
     val caught = intercept[SparkException] {
       relation.parseRecords(testReader, oneRowRDD).collect()
@@ -123,6 +126,51 @@ class CobolRelationSpec extends SparkCobolTestBase with Serializable {
     assert(caught.getMessage.contains("key not found: absentField"))
   }
 
+  it should "apply an exact global record limit" in {
+    val records = sqlContext.sparkContext.parallelize(
+      Seq(Row(1), Row(2), Row(3), Row(4), Row(5), Row(6)), 3)
+
+    val limited = CobolRelation.applyRecordLimit(records, Some(3))
+
+    limited.getNumPartitions shouldBe 1
+    limited.collect().toSeq shouldBe Seq(Row(1), Row(2), Row(3))
+  }
+
+  it should "continue to a later partition to satisfy a record limit" in {
+    val records = sqlContext.sparkContext.parallelize(Seq(Row(1), Row(2), Row(3)), 3)
+
+    val limited = CobolRelation.applyRecordLimit(records, Some(2))
+
+    limited.collect().toSeq shouldBe Seq(Row(1), Row(2))
+  }
+
+  it should "preserve records unchanged when a record limit is absent" in {
+    val records = sqlContext.sparkContext.parallelize(Seq(Row(1), Row(2), Row(3)), 3)
+
+    val limited = CobolRelation.applyRecordLimit(records, None)
+
+    limited should be theSameInstanceAs records
+    limited.getNumPartitions shouldBe 3
+    limited.collect().toSet shouldBe Set(Row(1), Row(2), Row(3))
+  }
+
+  it should "return an empty RDD without evaluating parent partitions for a zero record limit" in {
+    val records = sqlContext.sparkContext.parallelize(Seq(1), 1).mapPartitions[Row] { _ =>
+      throw new RuntimeException("Parent partition should not be evaluated")
+    }
+
+    CobolRelation.applyRecordLimit(records, Some(0)).collect().isEmpty shouldBe true
+  }
+
+  it should "not evaluate later parent partitions once a record limit is met" in {
+    val records = sqlContext.sparkContext.parallelize(Seq(1, 2), 2).mapPartitionsWithIndex {
+      case (0, _) => Iterator(Row(1), Row(2))
+      case (_, _) => throw new RuntimeException("Later parent partition should not be evaluated")
+    }
+
+    CobolRelation.applyRecordLimit(records, Some(2)).collect().toSeq shouldBe Seq(Row(1), Row(2))
+  }
+
   private def createTestData(): List[Map[String, Option[String]]] = {
     List(
       Map("name" -> Some("Young Guy"), "id" -> Some("1"), "age" -> Some("20")),
@@ -134,4 +182,4 @@ class CobolRelationSpec extends SparkCobolTestBase with Serializable {
   private def createTestSchemaMap(): Map[String, Any] = {
     Map("name" -> "", "id" -> "1l", "age" -> "1")
   }
-}
+}