From bedc4076f3505e750dec7efaa3923086dd736c4e Mon Sep 17 00:00:00 2001 From: Ruslan Iushchenko Date: Mon, 27 Jul 2026 12:10:03 +0200 Subject: [PATCH 1/4] #857 Disable indexes when record_limit is below threshold and pass recordLimit to VRLRecordReader. --- README.md | 28 +++++++++---------- .../reader/iterator/VRLRecordReader.scala | 6 +++- .../iterator/VarLenHierarchicalIterator.scala | 2 +- .../iterator/VarLenNestedIterator.scala | 2 +- .../reader/parameters/CobolParameters.scala | 4 +-- .../parameters/CobolParametersParser.scala | 20 +++++++++++-- .../reader/parameters/ReaderParameters.scala | 4 ++- .../reader/iterator/VRLRecordReaderSpec.scala | 3 +- .../spark/cobol/SparkCobolProcessor.scala | 2 +- .../spark/cobol/source/DefaultSource.scala | 2 +- .../integration/Test43RecordLimitSpec.scala | 21 +++++++++++++- 11 files changed, 67 insertions(+), 27 deletions(-) diff --git a/README.md b/README.md index 7d9526e3..07ddc427 100644 --- a/README.md +++ b/README.md @@ -1675,20 +1675,20 @@ The output looks like this: ##### Modifier options -| Option (usage example) | Description | -|-----------------------------------------------------|:------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| .option("schema_retention_policy", "collapse_root") | When `collapse_root` (default) the root level record will be removed from the Spark schema. When `keep_original`, the root level GROUP will be present in the Spark schema | -| .option("drop_group_fillers", "false") | If `true`, all GROUP FILLERs will be dropped from the output schema. If `false` (default), such fields will be retained. | -| .option("drop_value_fillers", "false") | If `true` (default), all non-GROUP FILLERs will be dropped from the output schema. If `false`, such fields will be retained. | -| .option("filler_naming_policy", "sequence_numbers") | Filler renaming strategy so that column names are not duplicated. Either `sequence_numbers` (default) or `previous_field_name` can be used. | -| .option("non_terminals", "GROUP1,GROUP2") | Specifies groups to also be added to the schema as string fields. When this option is specified, the reader will add one extra data field after each matching group containing the string data for the group. | -| .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. | +| Option (usage example) | Description | +|-----------------------------------------------------|:------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| .option("schema_retention_policy", "collapse_root") | When `collapse_root` (default) the root level record will be removed from the Spark schema. When `keep_original`, the root level GROUP will be present in the Spark schema | +| .option("drop_group_fillers", "false") | If `true`, all GROUP FILLERs will be dropped from the output schema. If `false` (default), such fields will be retained. | +| .option("drop_value_fillers", "false") | If `true` (default), all non-GROUP FILLERs will be dropped from the output schema. If `false`, such fields will be retained. | +| .option("filler_naming_policy", "sequence_numbers") | Filler renaming strategy so that column names are not duplicated. Either `sequence_numbers` (default) or `previous_field_name` can be used. | +| .option("non_terminals", "GROUP1,GROUP2") | Specifies groups to also be added to the schema as string fields. When this option is specified, the reader will add one extra data field after each matching group containing the string data for the group. | +| .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: diff --git a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VRLRecordReader.scala b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VRLRecordReader.scala index 88a00bb6..c5c6c5a1 100644 --- a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VRLRecordReader.scala +++ b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VRLRecordReader.scala @@ -35,6 +35,7 @@ import za.co.absa.cobrix.cobol.reader.validator.ReaderParametersValidator * @param recordExtractor A record extractor that can be used instead of the record header parser. * @param startRecordId A starting record id value for this particular file/stream `dataStream` * @param startingFileOffset An offset of the file where parsing should be started + * @param recordLimit An optional limit on the number of records to read */ class VRLRecordReader(cobolSchema: Copybook, dataStream: SimpleStream, @@ -42,7 +43,8 @@ class VRLRecordReader(cobolSchema: Copybook, recordHeaderParser: RecordHeaderParser, recordExtractor: Option[RawRecordExtractor], startRecordId: Long, - startingFileOffset: Long) extends Iterator[(String, Array[Byte])] with Logging { + startingFileOffset: Long, + recordLimit: Option[Int]) extends Iterator[(String, Array[Byte])] with Logging { type RawRecord = (String, Array[Byte]) @@ -72,6 +74,8 @@ class VRLRecordReader(cobolSchema: Copybook, @throws(classOf[IllegalStateException]) private def fetchNext(): Unit = { + if (recordLimit.isDefined && recordLimit.get <= recordIndex) + return var recordFetched = false while (!recordFetched) { val binaryData = recordExtractor match { diff --git a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VarLenHierarchicalIterator.scala b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VarLenHierarchicalIterator.scala index 1d413a5c..65d59f6c 100644 --- a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VarLenHierarchicalIterator.scala +++ b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VarLenHierarchicalIterator.scala @@ -54,7 +54,7 @@ final class VarLenHierarchicalIterator[T: ClassTag](cobolSchema: Copybook, type RawData = Array[Byte] type RawRecord = (String, Array[Byte]) - private val rawRecordIterator = new VRLRecordReader(cobolSchema, dataStream, readerProperties, recordHeaderParser, rawRecordExtractor, startRecordId, startingFileOffset) + private val rawRecordIterator = new VRLRecordReader(cobolSchema, dataStream, readerProperties, recordHeaderParser, rawRecordExtractor, startRecordId, startingFileOffset, readerProperties.recordLimit) private var recordIndex = startRecordId private var cachedValue: Option[Seq[Any]] = _ diff --git a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VarLenNestedIterator.scala b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VarLenNestedIterator.scala index a9bb532b..ebe991b6 100644 --- a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VarLenNestedIterator.scala +++ b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VarLenNestedIterator.scala @@ -52,7 +52,7 @@ final class VarLenNestedIterator[T: ClassTag](cobolSchema: Copybook, segmentIdPrefix: String, handler: RecordHandler[T]) extends Iterator[Seq[Any]] { - private val rawRecordIterator = new VRLRecordReader(cobolSchema, dataStream, readerProperties, recordHeaderParser, recordExtractor, startRecordId, startingFileOffset) + private val rawRecordIterator = new VRLRecordReader(cobolSchema, dataStream, readerProperties, recordHeaderParser, recordExtractor, startRecordId, startingFileOffset, readerProperties.recordLimit) private var cachedValue: Option[Seq[Any]] = _ private val segmentIdFilter = readerProperties.multisegment.flatMap(p => p.segmentIdFilter) diff --git a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParameters.scala b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParameters.scala index eb34a0e7..012639d4 100644 --- a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParameters.scala +++ b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParameters.scala @@ -117,7 +117,7 @@ case class CobolParameters( metadataPolicy: MetadataPolicy, fileHeaderField: Option[String] = None, fileTrailerField: Option[String] = None, + recordLimit: Option[Int] = None, writerParameters: Option[WriterParameters], - options: Map[String, String], - recordLimit: Option[Int] = None + options: Map[String, String] ) 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 31d6dd2a..c39e19cc 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,8 @@ object CobolParametersParser extends Logging { 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 MIN_RECORDS_FOR_INDEXES = 100000 + private def getSchemaRetentionPolicy(params: Parameters): SchemaRetentionPolicy = { val schemaRetentionPolicyName = params.getOrElse(PARAM_SCHEMA_RETENTION_POLICY, "collapse_root") val schemaRetentionPolicy = SchemaRetentionPolicy.withNameOpt(schemaRetentionPolicyName) @@ -354,9 +356,9 @@ object CobolParametersParser extends Logging { MetadataPolicy(params.getOrElse(PARAM_METADATA, "basic")), 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, writerParameters, - params.getMap, - recordLimit + params.getMap ) validateSparkCobolOptions(params, recordFormat, validateRedundantOptions) cobolParameters @@ -537,6 +539,7 @@ object CobolParametersParser extends Logging { varLenParams.inputFileNameColumn, parameters.metadataPolicy, recordsToExclude, + parameters.recordLimit, parameters.writerParameters, parameters.options ) @@ -564,6 +567,17 @@ object CobolParametersParser extends Logging { throw new IllegalArgumentException(s"Option '$PARAM_RECORD_LENGTH_FIELD' cannot be used together with '$PARAM_IS_RECORD_SEQUENCE' or '$PARAM_IS_XCOM'.") } + val recordLimit = getRecordLimit(params: Parameters) + val enableIndexes = recordLimit match { + case Some(limit) if limit > MIN_RECORDS_FOR_INDEXES => + params.getOrElse(PARAM_ENABLE_INDEXES, "true").toBoolean + case Some(limit) => + logger.info(s"Indexes are disabled for $PARAM_RECORD_LIMIT=$limit") + false + case None => + params.getOrElse(PARAM_ENABLE_INDEXES, "true").toBoolean + } + if (recordLengthFieldOpt.isDefined || isRecordSequence || isRecordIdGenerationEnabled || @@ -591,7 +605,7 @@ object CobolParametersParser extends Logging { fileStartOffset, fileEndOffset, isRecordIdGenerationEnabled, - params.getOrElse(PARAM_ENABLE_INDEXES, "true").toBoolean, + enableIndexes, params.getOrElse(PARAM_ENABLE_INDEX_CACHE, "true").toBoolean, params.get(PARAM_INPUT_SPLIT_RECORDS).map(v => v.toInt), params.get(PARAM_INPUT_SPLIT_SIZE_MB).map(v => v.toInt), diff --git a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/ReaderParameters.scala b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/ReaderParameters.scala index 4f486cdd..75b39056 100644 --- a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/ReaderParameters.scala +++ b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/ReaderParameters.scala @@ -81,6 +81,7 @@ import za.co.absa.cobrix.cobol.reader.policies.SchemaRetentionPolicy.SchemaReten * @param inputFileNameColumn A column name to add to the dataframe. The column will contain input file name for each record similar to 'input_file_name()' function * @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 writerParameters If specified, contains parameters for the writer. * @param options Options passed to spark-cobol */ @@ -143,7 +144,8 @@ case class ReaderParameters( reAdditionalInfo: String = "", inputFileNameColumn: String = "", metadataPolicy: MetadataPolicy = MetadataPolicy.Basic, - recordsToExclude: Set[String] = Set.empty, + recordsToExclude: Set[String] = Set.empty, + recordLimit: Option[Int] = None, writerParameters: Option[WriterParameters] = None, options: Map[String, String] = Map.empty ) diff --git a/cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/reader/iterator/VRLRecordReaderSpec.scala b/cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/reader/iterator/VRLRecordReaderSpec.scala index 7b1dda34..975a3e31 100644 --- a/cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/reader/iterator/VRLRecordReaderSpec.scala +++ b/cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/reader/iterator/VRLRecordReaderSpec.scala @@ -352,7 +352,8 @@ class VRLRecordReaderSpec extends AnyWordSpec { recordHeaderParser, recordExtractor, 0, - 0 + 0, + None ) } diff --git a/spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/SparkCobolProcessor.scala b/spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/SparkCobolProcessor.scala index 9f2aed98..2a78574e 100644 --- a/spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/SparkCobolProcessor.scala +++ b/spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/SparkCobolProcessor.scala @@ -207,7 +207,7 @@ object SparkCobolProcessor { // This method might be called several times during the ebcdic file processing. If there are redundant options, they will be logged. // `ignoreRedundantOptions=true` when the method is called just for copybook parsing so logging redundant options could be skipped. CobolParametersParser.parse(new Parameters(varLenOptions), !ignoreRedundantOptions) - .copy(sourcePaths = listOfFiles, copybookContent = Option(copybookContents)) + .copy(copybookContent = Option(copybookContents), sourcePaths = listOfFiles) } private def getRecordRdd(listOfFiles: Seq[String], diff --git a/spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/source/DefaultSource.scala b/spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/source/DefaultSource.scala index 8d9ac5c8..80ff26d6 100644 --- a/spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/source/DefaultSource.scala +++ b/spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/source/DefaultSource.scala @@ -283,7 +283,7 @@ object DefaultSource { case other => other } - parameters.copy(variableLengthParams = updatedVarLen, recordLength = updatedRecordLength) + parameters.copy(recordLength = updatedRecordLength, variableLengthParams = updatedVarLen) } /** diff --git a/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/source/integration/Test43RecordLimitSpec.scala b/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/source/integration/Test43RecordLimitSpec.scala index 23b6db32..dd711994 100644 --- a/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/source/integration/Test43RecordLimitSpec.scala +++ b/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/source/integration/Test43RecordLimitSpec.scala @@ -34,7 +34,10 @@ class Test43RecordLimitSpec extends AnyWordSpec with SparkTestBase with BinaryFi private val rdwData = Array( 0x00, 0x02, 0x00, 0x00, 0x30, 0x31, 0x00, 0x02, 0x00, 0x00, 0x32, 0x33, - 0x00, 0x02, 0x00, 0x00, 0x34, 0x35 + 0x00, 0x02, 0x00, 0x00, 0x34, 0x35, + 0x00, 0x02, 0x00, 0x00, 0x36, 0x37, + 0x00, 0x02, 0x00, 0x00, 0x38, 0x39, + 0x00, 0x02, 0x00, 0x00, 0x3A, 0x3B ).map(_.toByte) "record_limit" should { @@ -98,6 +101,22 @@ class Test43RecordLimitSpec extends AnyWordSpec with SparkTestBase with BinaryFi } } + "limit variable-length RDW records which allows indexing" in { + withTempBinFile("record_limit_rdw", ".dat", rdwData) { path => + val df = spark + .read + .format("cobol") + .option("copybook_contents", fixedLengthCopybook) + .option("encoding", "ascii") + .option("record_format", "V") + .option("is_rdw_big_endian", "true") + .option("record_limit", "1000000") + .load(path) + + assert(df.count() == 6) + } + } + "be accepted in pedantic mode" in { withTempBinFile("record_limit_pedantic", ".dat", "00010203".getBytes(StandardCharsets.US_ASCII)) { path => val df = fixedLengthDataFrame(Seq(path), Map("pedantic" -> "true", "record_limit" -> "1")) From f7aedb713a110ce7f27ee077bcb96c871ff14789 Mon Sep 17 00:00:00 2001 From: Ruslan Iushchenko Date: Mon, 27 Jul 2026 12:47:48 +0200 Subject: [PATCH 2/4] Fix record limit handling and stream closing in VRLRecordReader. --- .../cobrix/cobol/reader/iterator/VRLRecordReader.scala | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VRLRecordReader.scala b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VRLRecordReader.scala index c5c6c5a1..bb1e9e29 100644 --- a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VRLRecordReader.scala +++ b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VRLRecordReader.scala @@ -51,6 +51,7 @@ class VRLRecordReader(cobolSchema: Copybook, private var cachedValue: Option[RawRecord] = _ private var byteIndex = startingFileOffset private var recordIndex = startRecordId - 1 + private var recordNextToFetch = startRecordId final private val segmentIdField = ReaderParametersValidator.getSegmentIdField(readerProperties.multisegment, cobolSchema) final private val minimumRecordLength = readerProperties.minimumRecordLength @@ -74,13 +75,16 @@ class VRLRecordReader(cobolSchema: Copybook, @throws(classOf[IllegalStateException]) private def fetchNext(): Unit = { - if (recordLimit.isDefined && recordLimit.get <= recordIndex) + if (recordLimit.isDefined && recordLimit.get <= recordNextToFetch) { + dataStream.close() return + } var recordFetched = false while (!recordFetched) { val binaryData = recordExtractor match { case Some(extractor) => if (extractor.hasNext) { + recordNextToFetch += 1 Option(extractor.next()) } else { None @@ -96,6 +100,7 @@ class VRLRecordReader(cobolSchema: Copybook, case Some(data) if data.length < minimumRecordLength || data.length > maximumRecordLength => recordFetched = false case Some(data) => + recordNextToFetch += 1 val segmentId = getSegmentId(data) val segmentIdStr = segmentId.getOrElse("") From 673e509e53b8503a304c5010e5ae4f75758a105c Mon Sep 17 00:00:00 2001 From: Ruslan Iushchenko Date: Mon, 27 Jul 2026 13:56:39 +0200 Subject: [PATCH 3/4] #857 Fix record limit handling for hierarchical files. --- .../reader/iterator/VRLRecordReader.scala | 3 +- .../iterator/VarLenHierarchicalIterator.scala | 10 ++++- .../integration/Test43RecordLimitSpec.scala | 37 +++++++++++++++++++ 3 files changed, 47 insertions(+), 3 deletions(-) diff --git a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VRLRecordReader.scala b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VRLRecordReader.scala index bb1e9e29..781103a1 100644 --- a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VRLRecordReader.scala +++ b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VRLRecordReader.scala @@ -84,7 +84,6 @@ class VRLRecordReader(cobolSchema: Copybook, val binaryData = recordExtractor match { case Some(extractor) => if (extractor.hasNext) { - recordNextToFetch += 1 Option(extractor.next()) } else { None @@ -100,10 +99,10 @@ class VRLRecordReader(cobolSchema: Copybook, case Some(data) if data.length < minimumRecordLength || data.length > maximumRecordLength => recordFetched = false case Some(data) => - recordNextToFetch += 1 val segmentId = getSegmentId(data) val segmentIdStr = segmentId.getOrElse("") + recordNextToFetch += 1 cachedValue = Some(segmentIdStr, data) recordFetched = true } diff --git a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VarLenHierarchicalIterator.scala b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VarLenHierarchicalIterator.scala index 65d59f6c..ed74d1fe 100644 --- a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VarLenHierarchicalIterator.scala +++ b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VarLenHierarchicalIterator.scala @@ -54,9 +54,10 @@ final class VarLenHierarchicalIterator[T: ClassTag](cobolSchema: Copybook, type RawData = Array[Byte] type RawRecord = (String, Array[Byte]) - private val rawRecordIterator = new VRLRecordReader(cobolSchema, dataStream, readerProperties, recordHeaderParser, rawRecordExtractor, startRecordId, startingFileOffset, readerProperties.recordLimit) + private val rawRecordIterator = new VRLRecordReader(cobolSchema, dataStream, readerProperties, recordHeaderParser, rawRecordExtractor, startRecordId, startingFileOffset, None) private var recordIndex = startRecordId + private var rootRecordIndex = startRecordId private var cachedValue: Option[Seq[Any]] = _ private val segmentRedefines = cobolSchema.getAllSegmentRedefines.toArray @@ -98,6 +99,11 @@ final class VarLenHierarchicalIterator[T: ClassTag](cobolSchema: Copybook, @throws(classOf[IllegalStateException]) private def fetchNext(): Unit = { + if (readerProperties.recordLimit.isDefined && readerProperties.recordLimit.get <= rootRecordIndex) { + dataStream.close() + return + } + var recordFetched = false while (!recordFetched) { if (rawRecordIterator.hasNext) { @@ -134,6 +140,8 @@ final class VarLenHierarchicalIterator[T: ClassTag](cobolSchema: Copybook, recordFetched = true } } + if (recordFetched) + rootRecordIndex = rootRecordIndex + 1 } private def extractRow(records: ArrayBuffer[RawRecord]): Seq[Any] = { diff --git a/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/source/integration/Test43RecordLimitSpec.scala b/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/source/integration/Test43RecordLimitSpec.scala index dd711994..8d060456 100644 --- a/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/source/integration/Test43RecordLimitSpec.scala +++ b/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/source/integration/Test43RecordLimitSpec.scala @@ -117,6 +117,43 @@ class Test43RecordLimitSpec extends AnyWordSpec with SparkTestBase with BinaryFi } } + "limit of hierarchical data should work" in { + val rdwHierarchicalData = Array( + 0x00, 0x02, 0x00, 0x00, 0xF0, 0x81, + 0x00, 0x02, 0x00, 0x00, 0xF1, 0xC2, + 0x00, 0x02, 0x00, 0x00, 0xF1, 0xC3, + 0x00, 0x02, 0x00, 0x00, 0xF0, 0x84, + 0x00, 0x02, 0x00, 0x00, 0xF0, 0x85, + 0x00, 0x02, 0x00, 0x00, 0xF1, 0xC6 + ).map(_.toByte) + + val copybook = + """ 01 R. + | 03 SEGMENT-ID PIC 9(1). + | 03 SEGMENT0. + | 10 A PIC X(1). + | 03 SEGMENT1 REDEFINES SEGMENT0. + | 10 B PIC X(1). + |""".stripMargin + + withTempBinFile("record_limit_rdw", ".dat", rdwHierarchicalData) { path => + val df = spark + .read + .format("cobol") + .option("copybook_contents", copybook) + .option("record_format", "V") + .option("is_rdw_big_endian", "true") + .option("record_limit", "2") + .option("segment_field", "SEGMENT_ID") + .option("redefine_segment_id_map:1", "SEGMENT0 => 0") + .option("redefine_segment_id_map:2", "SEGMENT1 => 1") + .option("segment-children:1", "SEGMENT0 => SEGMENT1") + .load(path) + + assert(df.count() == 2) + } + } + "be accepted in pedantic mode" in { withTempBinFile("record_limit_pedantic", ".dat", "00010203".getBytes(StandardCharsets.US_ASCII)) { path => val df = fixedLengthDataFrame(Seq(path), Map("pedantic" -> "true", "record_limit" -> "1")) From f971d8701417fa6c6baee61a0044a83894a20155 Mon Sep 17 00:00:00 2001 From: Ruslan Iushchenko Date: Mon, 27 Jul 2026 14:09:11 +0200 Subject: [PATCH 4/4] Fix record limit to clear cached value before closing stream. --- .../co/absa/cobrix/cobol/reader/iterator/VRLRecordReader.scala | 1 + .../cobol/reader/iterator/VarLenHierarchicalIterator.scala | 1 + 2 files changed, 2 insertions(+) diff --git a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VRLRecordReader.scala b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VRLRecordReader.scala index 781103a1..f889f14e 100644 --- a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VRLRecordReader.scala +++ b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VRLRecordReader.scala @@ -76,6 +76,7 @@ class VRLRecordReader(cobolSchema: Copybook, @throws(classOf[IllegalStateException]) private def fetchNext(): Unit = { if (recordLimit.isDefined && recordLimit.get <= recordNextToFetch) { + cachedValue = None dataStream.close() return } diff --git a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VarLenHierarchicalIterator.scala b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VarLenHierarchicalIterator.scala index ed74d1fe..235d21c1 100644 --- a/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VarLenHierarchicalIterator.scala +++ b/cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VarLenHierarchicalIterator.scala @@ -100,6 +100,7 @@ final class VarLenHierarchicalIterator[T: ClassTag](cobolSchema: Copybook, @throws(classOf[IllegalStateException]) private def fetchNext(): Unit = { if (readerProperties.recordLimit.isDefined && readerProperties.recordLimit.get <= rootRecordIndex) { + cachedValue = None dataStream.close() return }