Skip to content

#857 Disable indexes when record_limit is below threshold and pass recordLimit to VRLRecordReader. #861

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yruslan merged 4 commits into from
Jul 27, 2026
Merged

#857 Disable indexes when record_limit is below threshold and pass recordLimit to VRLRecordReader. #861

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
bedc407
#857 Disable indexes when record_limit is below threshold and pass re…
yruslan Jul 27, 2026
f7aedb7
Fix record limit handling and stream closing in VRLRecordReader.
yruslan Jul 27, 2026
673e509
#857 Fix record limit handling for hierarchical files.
yruslan Jul 27, 2026
f971d87
Fix record limit to clear cached value before closing stream.
yruslan Jul 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 14 additions & 14 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:

Expand Down
11 changes: 10 additions & 1 deletion cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VRLRecordReader.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,20 +35,23 @@ 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,
readerProperties: ReaderParameters,
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])

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
Expand All @@ -72,6 +75,11 @@ class VRLRecordReader(cobolSchema: Copybook,

@throws(classOf[IllegalStateException])
private def fetchNext(): Unit = {
if (recordLimit.isDefined && recordLimit.get <= recordNextToFetch) {
cachedValue = None
dataStream.close()
return
}
var recordFetched = false
while (!recordFetched) {
val binaryData = recordExtractor match {
Expand All @@ -95,6 +103,7 @@ class VRLRecordReader(cobolSchema: Copybook,
val segmentId = getSegmentId(data)
val segmentIdStr = segmentId.getOrElse("")

recordNextToFetch += 1
cachedValue = Some(segmentIdStr, data)
recordFetched = true
}
Expand Down
11 changes: 10 additions & 1 deletion ...r/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VarLenHierarchicalIterator.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
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

Expand Down Expand Up @@ -98,6 +99,12 @@ 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
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.

var recordFetched = false
while (!recordFetched) {
if (rawRecordIterator.hasNext) {
Expand Down Expand Up @@ -134,6 +141,8 @@ final class VarLenHierarchicalIterator[T: ClassTag](cobolSchema: Copybook,
recordFetched = true
}
}
if (recordFetched)
rootRecordIndex = rootRecordIndex + 1
}

private def extractRow(records: ArrayBuffer[RawRecord]): Seq[Any] = {
Expand Down
2 changes: 1 addition & 1 deletion ...-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/iterator/VarLenNestedIterator.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
4 changes: 2 additions & 2 deletions cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParameters.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
)
20 changes: 17 additions & 3 deletions ...rser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParametersParser.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -537,6 +539,7 @@ object CobolParametersParser extends Logging {
varLenParams.inputFileNameColumn,
parameters.metadataPolicy,
recordsToExclude,
parameters.recordLimit,
parameters.writerParameters,
parameters.options
)
Expand Down Expand Up @@ -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 ||
Expand Down Expand Up @@ -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),
Expand Down
4 changes: 3 additions & 1 deletion cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/parameters/ReaderParameters.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
*/
Expand Down Expand Up @@ -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
)
3 changes: 2 additions & 1 deletion ...l-parser/src/test/scala/za/co/absa/cobrix/cobol/reader/iterator/VRLRecordReaderSpec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -352,7 +352,8 @@ class VRLRecordReaderSpec extends AnyWordSpec {
recordHeaderParser,
recordExtractor,
0,
0
0,
None
)
}

Expand Down
Loading
Loading