diff --git a/README.md b/README.md index 48b03c4d..7d9526e3 100644 --- a/README.md +++ b/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 | 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 a716a00a..eb34a0e7 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 @@ -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 ) 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 1b0fbced..31d6dd2a 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 @@ -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 diff --git a/cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParametersParserSuite.scala b/cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParametersParserSuite.scala index 5257f22d..265ae423 100644 --- a/cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/reader/parameters/CobolParametersParserSuite.scala +++ b/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", diff --git a/spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/source/CobolRelation.scala b/spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/source/CobolRelation.scala index 55339c23..9b3c8fcd 100644 --- a/spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/source/CobolRelation.scala +++ b/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) } } -} \ No newline at end of file +} 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 60a8cbbe..8d9ac5c8 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 @@ -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 */ diff --git a/spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/source/parameters/CobolParametersValidator.scala b/spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/source/parameters/CobolParametersValidator.scala index 777a0c67..ab377627 100644 --- a/spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/source/parameters/CobolParametersValidator.scala +++ b/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("; ")}") } diff --git a/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/source/CobolRelationSpec.scala b/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/source/CobolRelationSpec.scala index e233488b..25254df1 100644 --- a/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/source/CobolRelationSpec.scala +++ b/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") } -} \ No newline at end of file +} 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 new file mode 100644 index 00000000..23b6db32 --- /dev/null +++ b/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/source/integration/Test43RecordLimitSpec.scala @@ -0,0 +1,124 @@ +/* + * 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.spark.cobol.source.integration + +import org.apache.spark.sql.DataFrame +import org.scalatest.wordspec.AnyWordSpec +import za.co.absa.cobrix.spark.cobol.source.base.SparkTestBase +import za.co.absa.cobrix.spark.cobol.source.fixtures.BinaryFileFixture + +import java.nio.charset.StandardCharsets + +class Test43RecordLimitSpec extends AnyWordSpec with SparkTestBase with BinaryFileFixture { + private val fixedLengthCopybook = + """ 01 R. + | 03 A PIC X(2). + |""".stripMargin + + private val fixedLengthValues = Seq("00", "01", "02", "03", "10", "11", "12", "13") + + private val rdwData = Array( + 0x00, 0x02, 0x00, 0x00, 0x30, 0x31, + 0x00, 0x02, 0x00, 0x00, 0x32, 0x33, + 0x00, 0x02, 0x00, 0x00, 0x34, 0x35 + ).map(_.toByte) + + "record_limit" should { + "apply globally across multiple fixed-length input paths" in { + withTempBinFile("record_limit_fixed_1", ".dat", "00010203".getBytes(StandardCharsets.US_ASCII)) { firstPath => + withTempBinFile("record_limit_fixed_2", ".dat", "10111213".getBytes(StandardCharsets.US_ASCII)) { secondPath => + val actual = values(fixedLengthDataFrame(Seq(firstPath, secondPath), Map("record_limit" -> "3"))) + + assert(actual.length == 3) + assert(actual.distinct.length == actual.length) + assert(actual.forall(fixedLengthValues.contains)) + } + } + } + + "return no rows when set to zero" in { + withTempBinFile("record_limit_zero", ".dat", "00010203".getBytes(StandardCharsets.US_ASCII)) { path => + assert(fixedLengthDataFrame(Seq(path), Map("record_limit" -> "0")).count() == 0) + } + } + + "leave all rows available when absent" in { + withTempBinFile("record_limit_absent_1", ".dat", "00010203".getBytes(StandardCharsets.US_ASCII)) { firstPath => + withTempBinFile("record_limit_absent_2", ".dat", "10111213".getBytes(StandardCharsets.US_ASCII)) { secondPath => + val actual = values(fixedLengthDataFrame(Seq(firstPath, secondPath))) + + assert(actual.length == fixedLengthValues.length) + assert(actual.toSet == fixedLengthValues.toSet) + } + } + } + + "limit UTF-8 text records in D format" in { + withTempTextFile("record_limit_text", ".txt", StandardCharsets.UTF_8, "aa\nbb\ncc") { path => + val df = spark + .read + .format("cobol") + .option("copybook_contents", fixedLengthCopybook) + .option("record_format", "D") + .option("ascii_charset", "UTF-8") + .option("record_limit", "2") + .load(path) + + assert(values(df) == Seq("aa", "bb")) + } + } + + "limit variable-length RDW records" 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", "2") + .load(path) + + assert(values(df) == Seq("01", "23")) + } + } + + "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")) + + assert(values(df) == Seq("00")) + } + } + } + + private def fixedLengthDataFrame(paths: Seq[String], options: Map[String, String] = Map.empty): DataFrame = { + spark + .read + .format("cobol") + .option("copybook_contents", fixedLengthCopybook) + .option("encoding", "ascii") + .option("data_paths", paths.mkString(",")) + .options(options) + .load() + } + + private def values(df: DataFrame): Seq[String] = { + df.select("A").collect().map(_.getString(0)).toSeq + } +} diff --git a/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/source/parameters/CobolParametersValidatorSuite.scala b/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/source/parameters/CobolParametersValidatorSuite.scala index 973fd2da..54ae9a2f 100644 --- a/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/source/parameters/CobolParametersValidatorSuite.scala +++ b/spark-cobol/src/test/scala/za/co/absa/cobrix/spark/cobol/source/parameters/CobolParametersValidatorSuite.scala @@ -29,7 +29,8 @@ class CobolParametersValidatorSuite extends AnyWordSpec { occursMappings = Map("A" -> Map("B" -> 1)), startOffset = 1, fileEndOffset = 2, - multisegment = Some(MultisegmentParameters("SEG", None, Seq.empty, "", Map.empty, Map.empty)) + multisegment = Some(MultisegmentParameters("SEG", None, Seq.empty, "", Map.empty, Map.empty)), + options = Map("record_limit" -> "10") ) val ex = intercept[IllegalArgumentException] { @@ -41,6 +42,7 @@ class CobolParametersValidatorSuite extends AnyWordSpec { assert(ex.getMessage.contains("'record_start_offset' and 'record_end_offset' are not supported for writing")) assert(ex.getMessage.contains("'file_start_offset' and 'file_end_offset' are not supported for writing")) assert(ex.getMessage.contains("Multi-segment options ('segment_field', 'segment_filter', etc) are not supported for writing")) + assert(ex.getMessage.contains("'record_limit' is supported only for reading")) } "do not throw exceptions if the configuration is okay" in {