Feature/fixed width parsing

.gitignore

@@ -5,6 +5,7 @@
 .cache
 .history
 .lib/
+.idea/
 dist/*
 target/
 lib_managed/

README.md

@@ -1,2 +1,78 @@
 # spark-fixedwidth
-Fixed-width data source for Spark SQL and DataFrames
+Fixed-width data source for Spark SQL and DataFrames. Based on (and uses) [databricks-spark-csv](https://github.com/databricks/spark-csv)
+
+## Requirements
+This library requires Spark 1.3+ and Scala 2.11+
+
+## Building
+Run `sbt assembly` from inside the root directory to generate a JAR
+
+## Running / Using
+
+### In the Spark Shell
+`./bin/spark-shell --jars <PATH_TO>/spark-fixedwidth/target/scala-2.11/spark-fixedwidth-assembly-1.0.jar`
+
+### In another project
+Add the JAR to your project lib and sbt will include it for you
+
+## Features
+This package allows reading fixed-width files in local or distributed filesystem as [Spark DataFrames](https://spark.apache.org/docs/1.3.0/sql-programming-guide.html).
+When reading files the API accepts several options:
+* `path` (REQUIRED): location of files. Similar to Spark can accept standard Hadoop globbing expressions.
+* `fixedWidths` (REQUIRED): Int array of the fixed widths of the source file(s)
+* `schema`: in [spark SQL form](http://spark.apache.org/docs/latest/api/scala/index.html#org.apache.spark.sql.types.StructType). Otherwise everything is assumed String (unless inferSchema is on)
+* `useHeader`: when set to true the first line of files will be used to name columns and will not be included in data. All types will be assumed string. Default value is true.
+* `charset`: defaults to 'UTF-8' but can be set to other valid charset names
+* `inferSchema`: automatically infers column types. It requires one extra pass over the data and is false by default
+* `comment`: skip lines beginning with this character. Default is `"#"`. Disable comments by setting this to `null`.
+* `mode`: determines the parsing mode. By default it is PERMISSIVE. Possible values are:
+  * `PERMISSIVE`: tries to parse all lines: nulls are inserted for missing tokens and extra tokens are ignored.
+  * `DROPMALFORMED`: drops lines which have fewer or more tokens than expected or tokens which do not match the schema
+  * `FAILFAST`: aborts with a RuntimeException if encounters any malformed line
+* `codec`: compression codec to use when saving to file. Should be the fully qualified name of a class implementing `org.apache.hadoop.io.compress.CompressionCodec` or one of case-insensitive shorten names (`bzip2`, `gzip`, `lz4`, and `snappy`). Defaults to no compression when a codec is not specified.
+* `nullValue`: specify a string that indicates a null value. Any fields matching this string will be set as nulls in the DataFrame
+* `ignoreLeadingWhiteSpace`: Boolean, default true
+* `ignoreTrailingWhiteSpace`: Boolean, default true
+
+### Scala API
+__Spark 1.4+:__
+
+See [sample fixed-width files](src/test/resources)
+```scala
+import org.apache.spark.sql.SQLContext
+import org.apache.spark.sql.types.{StructType, StructField, StringType, IntegerType, DoubleType}
+import com.quartethealth.spark.fixedwidth.FixedwidthContext
+
+val fruitSchema = StructType(Seq(
+    StructField("val", IntegerType),
+    StructField("name", StringType),
+    StructField("avail", StringType),
+    StructField("cost", DoubleType)
+))
+
+val sqlContext = new SQLContext(sc) // sc is defined in the spark console
+val fruitWidths = Array(3, 10, 5, 4)
+val fruit_resource = "fruit_fixedwidths.txt"
+
+val result = sqlContext.fixedFile(
+    fruit_resource,
+    fruitWidths,
+    fruitSchema,
+    useHeader = false
+)
+result.show() // Prints top 20 rows in tabular format
+
+// Example without schema, and showing extra options
+val fruit_resource = "fruit_w_headers_fixedwidths.txt"
+val result = sqlContext.fixedFile(
+    fruit_resource,
+    fruitWidths,
+    useHeader = true,
+    inferSchema = true,
+    mode = "DROPMALFORMED",
+    comment = '/',
+    ignoreLeadingWhiteSpace = true,
+    ignoreTrailingWhiteSpace = false
+)
+result.collect() // Returns an array that contains all of Rows in this DataFrame
+```

build.sbt

@@ -0,0 +1,18 @@
+name := "spark-fixedwidth"
+
+version := "1.0"
+
+organization := "com.quartethealth"
+
+scalaVersion := "2.11.7"
+
+libraryDependencies ++= Seq(
+  "com.univocity" % "univocity-parsers" % "1.5.1",
+  "org.apache.spark" %% "spark-sql" % "1.6.0" % "provided",
+  "org.apache.spark" %% "spark-core" % "1.6.0" % "provided",
+  "org.scalatest" %% "scalatest" % "2.2.1" % "test",
+
+  "org.specs2" %% "specs2-core" % "3.7" % "test"
+)
+
+scalacOptions in Test ++= Seq("-Yrangepos")

project/Build.scala

@@ -0,0 +1,8 @@
+import sbt._
+
+object MyBuild extends Build {
+
+  lazy val root = Project("root", file(".")) dependsOn csvProj
+  lazy val csvProj = RootProject(uri("git://github.com/quartethealth/spark-csv"))
+
+}

project/assembly.sbt

@@ -0,0 +1 @@
+addSbtPlugin("com.eed3si9n" % "sbt-assembly" % "0.13.0")

project/build.properties

@@ -0,0 +1 @@
+sbt.version = 0.13.8

project/plugins.sbt

@@ -0,0 +1 @@
+logLevel := Level.Warn

src/main/scala/com/quartethealth/spark/fixedwidth/FixedwidthRelation.scala

@@ -0,0 +1,56 @@
+package com.quartethealth.spark.fixedwidth
+
+import com.databricks.spark.csv.CsvRelation
+import com.databricks.spark.csv.readers.{BulkReader, LineReader}
+import com.quartethealth.spark.fixedwidth.readers.{BulkFixedwidthReader, LineFixedwidthReader}
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.SQLContext
+import org.apache.spark.sql.types.StructType
+
+class FixedwidthRelation protected[spark] (
+    baseRDD: () => RDD[String],
+    fixedWidths: Array[Int],
+    location: Option[String],
+    useHeader: Boolean,
+    parseMode: String,
+    comment: Character,
+    ignoreLeadingWhiteSpace: Boolean,
+    ignoreTrailingWhiteSpace: Boolean,
+    treatEmptyValuesAsNulls: Boolean,
+    userSchema: StructType,
+    inferSchema: Boolean,
+    codec: String = null,
+    nullValue: String = "")(@transient override val sqlContext: SQLContext)
+  extends CsvRelation(
+    baseRDD,
+    location,
+    useHeader,
+    delimiter = '\0',
+    quote = null,
+    escape = null,
+    comment = comment,
+    parseMode = parseMode,
+    parserLib = "UNIVOCITY",
+    ignoreLeadingWhiteSpace = ignoreLeadingWhiteSpace,
+    ignoreTrailingWhiteSpace = ignoreTrailingWhiteSpace,
+    treatEmptyValuesAsNulls = treatEmptyValuesAsNulls,
+    userSchema = userSchema,
+    inferCsvSchema = true,
+    codec = codec)(sqlContext) {
+
+  protected override def getLineReader(): LineReader = {
+    val commentChar: Char = if (comment == null) '\0' else comment
+    new LineFixedwidthReader(fixedWidths, commentMarker = commentChar,
+      ignoreLeadingSpace = ignoreLeadingWhiteSpace,
+      ignoreTrailingSpace = ignoreTrailingWhiteSpace)
+  }
+
+  protected override def getBulkReader(header: Seq[String], iter: Iterator[String],
+      split: Int): BulkReader = {
+    val commentChar: Char = if (comment == null) '\0' else comment
+    new BulkFixedwidthReader(iter, split, fixedWidths,
+      headers = header, commentMarker = commentChar,
+      ignoreLeadingSpace = ignoreLeadingWhiteSpace,
+      ignoreTrailingSpace = ignoreTrailingWhiteSpace)
+  }
+}

src/main/scala/com/quartethealth/spark/fixedwidth/package.scala

@@ -0,0 +1,35 @@
+package com.quartethealth.spark
+
+import com.databricks.spark.csv.util.TextFile
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.{DataFrame, SQLContext}
+
+package object fixedwidth {
+  implicit class FixedwidthContext(sqlContext: SQLContext) extends Serializable {
+    def fixedFile(
+        filePath: String,
+        fixedWidths: Array[Int],
+        schema: StructType = null,
+        useHeader: Boolean = true,
+        mode: String = "PERMISSIVE",
+        comment: Character = null,
+        ignoreLeadingWhiteSpace: Boolean = true,
+        ignoreTrailingWhiteSpace: Boolean = true,
+        charset: String = TextFile.DEFAULT_CHARSET.name(),
+        inferSchema: Boolean = false): DataFrame = {
+      val fixedwidthRelation = new FixedwidthRelation(
+        () => TextFile.withCharset(sqlContext.sparkContext, filePath, charset),
+        location = Some(filePath),
+        useHeader = useHeader,
+        comment = comment,
+        parseMode = mode,
+        fixedWidths = fixedWidths,
+        ignoreLeadingWhiteSpace = ignoreLeadingWhiteSpace,
+        ignoreTrailingWhiteSpace = ignoreTrailingWhiteSpace,
+        userSchema = schema,
+        inferSchema = inferSchema,
+        treatEmptyValuesAsNulls = false)(sqlContext)
+      sqlContext.baseRelationToDataFrame(fixedwidthRelation)
+    }
+  }
+}

src/main/scala/com/quartethealth/spark/fixedwidth/readers/readers.scala

@@ -0,0 +1,142 @@
+package com.quartethealth.spark.fixedwidth.readers
+
+import com.databricks.spark.csv.readers.{BulkReader, LineReader}
+import com.univocity.parsers.fixed.{FixedWidthFieldLengths, FixedWidthParser, FixedWidthParserSettings}
+
+/**
+  * Read and parse Fixed-width-like input
+  *
+  * @param fixedWidths the fixed widths of the fields
+  * @param lineSep the delimiter used to separate lines
+  * @param commentMarker Ignore lines starting with this char
+  * @param ignoreLeadingSpace ignore white space before a field
+  * @param ignoreTrailingSpace ignore white space after a field
+  * @param headers headers for the columns
+  * @param inputBufSize size of buffer to use for parsing input, tune for performance
+  * @param maxCols maximum number of columns allowed, for safety against bad inputs
+  */
+private[readers] abstract class FixedwidthReader(
+    fixedWidths: Array[Int],
+    lineSep: String = "\n",
+    commentMarker: Char = '#',
+    ignoreLeadingSpace: Boolean = true,
+    ignoreTrailingSpace: Boolean = true,
+    headers: Seq[String],
+    inputBufSize: Int = 128,
+    maxCols: Int = 20480) {
+  protected lazy val parser: FixedWidthParser = {
+    val settings = new FixedWidthParserSettings(new FixedWidthFieldLengths(fixedWidths: _*))
+    val format = settings.getFormat
+    format.setLineSeparator(lineSep)
+    format.setComment(commentMarker)
+    settings.setIgnoreLeadingWhitespaces(ignoreLeadingSpace)
+    settings.setIgnoreTrailingWhitespaces(ignoreTrailingSpace)
+    settings.setReadInputOnSeparateThread(false)
+    settings.setInputBufferSize(inputBufSize)
+    settings.setMaxColumns(maxCols)
+    settings.setNullValue("")
+    settings.setMaxCharsPerColumn(100000)
+    if (headers != null) settings.setHeaders(headers: _*)
+    // TODO: configurable?
+    settings.setSkipTrailingCharsUntilNewline(true)
+    settings.setRecordEndsOnNewline(true)
+
+    new FixedWidthParser(settings)
+  }
+}
+
+/**
+  * Read and parse a single line of Fixed-width-like input. Inefficient for bulk data.
+  * @param fixedWidths the fixed widths of the fields
+  * @param lineSep the delimiter used to separate lines
+  * @param commentMarker Ignore lines starting with this char
+  * @param ignoreLeadingSpace ignore white space before a field
+  * @param ignoreTrailingSpace ignore white space after a field
+  * @param inputBufSize size of buffer to use for parsing input, tune for performance
+  * @param maxCols maximum number of columns allowed, for safety against bad inputs
+  */
+private[fixedwidth] class LineFixedwidthReader(
+    fixedWidths: Array[Int],
+    lineSep: String = "\n",
+    commentMarker: Char = '#',
+    ignoreLeadingSpace: Boolean = true,
+    ignoreTrailingSpace: Boolean = true,
+    inputBufSize: Int = 128,
+    maxCols: Int = 20480)
+  extends FixedwidthReader(
+    fixedWidths,
+    lineSep,
+    commentMarker,
+    ignoreLeadingSpace,
+    ignoreTrailingSpace,
+    null,
+    inputBufSize,
+    maxCols)
+  with LineReader {
+  /**
+    * parse a line
+    * @param line a String with no newline at the end
+    * @return array of strings where each string is a field in the CSV record
+    */
+  def parseLine(line: String): Array[String] = {
+    parser.beginParsing(reader(line))
+    val parsed = parser.parseNext()
+    parser.stopParsing()
+    parsed
+  }
+}
+
+/**
+  * Read and parse Fixed-width-like input
+  *
+  * @param fixedWidths the fixed widths of the fields
+  * @param lineSep the delimiter used to separate lines
+  * @param commentMarker Ignore lines starting with this char
+  * @param ignoreLeadingSpace ignore white space before a field
+  * @param ignoreTrailingSpace ignore white space after a field
+  * @param headers headers for the columns
+  * @param inputBufSize size of buffer to use for parsing input, tune for performance
+  * @param maxCols maximum number of columns allowed, for safety against bad inputs
+  */
+private[fixedwidth] class BulkFixedwidthReader(
+    iter: Iterator[String],
+    split: Int,      // for debugging
+    fixedWidths: Array[Int],
+    lineSep: String = "\n",
+    commentMarker: Char = '#',
+    ignoreLeadingSpace: Boolean = true,
+    ignoreTrailingSpace: Boolean = true,
+    headers: Seq[String],
+    inputBufSize: Int = 128,
+    maxCols: Int = 20480)
+  extends FixedwidthReader(
+    fixedWidths,
+    lineSep,
+    commentMarker,
+    ignoreLeadingSpace,
+    ignoreTrailingSpace,
+    headers,
+    inputBufSize,
+    maxCols
+  ) with BulkReader {
+
+  parser.beginParsing(reader(iter))
+  private var nextRecord = parser.parseNext()
+
+  /**
+    * get the next parsed line.
+    *
+    * @return array of strings where each string is a field in the fixed-width record
+    */
+  override def next(): Array[String] = {
+    val curRecord = nextRecord
+    if(curRecord != null) {
+      nextRecord = parser.parseNext()
+    } else {
+      throw new NoSuchElementException("next record is null")
+    }
+    curRecord
+  }
+
+  override def hasNext: Boolean = nextRecord != null
+}

src/test/resources/fruit__fixedwidth.txt

@@ -0,0 +1,7 @@
+56 apple     TRUE 0.56
+45 pear      FALSE1.34
+34 raspberry TRUE 2.43
+34 plum      TRUE 1.31
+53 cherry    TRUE 1.4 
+23 orange    FALSE2.34
+56 persimmon FALSE23.2

src/test/resources/fruit_comments_fixedwidth.txt

@@ -0,0 +1,11 @@
+// Fruit for sale
+// 2014-10-02
+AMTNAME      AVAILCOST
+56 apple     TRUE 0.56
+45 pear      FALSE1.34
+34 raspberry TRUE 2.43
+34 plum      TRUE 1.31
+53 cherry    TRUE 1.4 
+23 orange    FALSE2.34
+56 persimmon FALSE23.2
+// No more fruit to show

src/test/resources/fruit_malformed_fixedwidth.txt

@@ -0,0 +1,7 @@
+23 apple     TRUE 1.44
+45 cherry    FALSE1.33
+34 raspberry TRUE BLR
+34
+ewfergerfgrefergregreg
+hi pear
+56 persimmon FALSE23.2