Skip to content

[SPARK-12871][SQL] Support to specify the option for compression codec. #10805

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

Closed
HyukjinKwon wants to merge 14 commits into from
Closed

[SPARK-12871][SQL] Support to specify the option for compression codec. #10805

Show file tree
Hide file tree
Changes from 10 commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
06774ad
Support for compress option
HyukjinKwon Jan 18, 2016
5e1611d
Correct Scala style and add an alias for codec as compression.
HyukjinKwon Jan 18, 2016
d154f02
Move back some codes changed unintentionally.
HyukjinKwon Jan 18, 2016
5b57fc2
Remove Hadoop 1.x configurations
HyukjinKwon Jan 18, 2016
e7ebddd
Look up compression first.
HyukjinKwon Jan 19, 2016
f4ffbf6
Resolve conflicts
HyukjinKwon Jan 19, 2016
adb9eb2
Support for shorten names and change the variable from codec to compr…
HyukjinKwon Jan 19, 2016
432da5d
Correct several nits and validate class names.
HyukjinKwon Jan 19, 2016
4388fe5
Add a newline between function and variables
HyukjinKwon Jan 19, 2016
6400b76
Remove whitespace
HyukjinKwon Jan 19, 2016
56316a8
Several nits and reduce complexity
HyukjinKwon Jan 19, 2016
c9217be
Update exception message
HyukjinKwon Jan 19, 2016
0245eea
Simply add a space after comma.
HyukjinKwon Jan 19, 2016
cd9f742
Update exception message and get rid of Option.
HyukjinKwon Jan 20, 2016
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
39 changes: 37 additions & 2 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/csv/CSVParameters.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,10 @@ package org.apache.spark.sql.execution.datasources.csv

import java.nio.charset.Charset

import org.apache.hadoop.io.compress._

import org.apache.spark.Logging
import org.apache.spark.util.Utils

private[sql] case class CSVParameters(@transient parameters: Map[String, String]) extends Logging {

Expand All @@ -35,7 +38,7 @@ private[sql] case class CSVParameters(@transient parameters: Map[String, String]

private def getBool(paramName: String, default: Boolean = false): Boolean = {
val param = parameters.getOrElse(paramName, default.toString)
if (param.toLowerCase() == "true") {
if (param.toLowerCase == "true") {
true
} else if (param.toLowerCase == "false") {
false
Expand Down Expand Up @@ -73,6 +76,14 @@ private[sql] case class CSVParameters(@transient parameters: Map[String, String]

val nullValue = parameters.getOrElse("nullValue", "")

val compressionCodecName =
Copy link
Contributor

@rxin rxin Jan 19, 2016

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

don't expose this variable. e.g. you can do

val compressionCodec: Option[String] = {
  val name = parameters.getOrElse("compression", parameters.get("codec"))
  name.map(CSVCompressionCodecs.getCodecClassName)
}

parameters.getOrElse("compression", parameters.getOrElse("codec", null))
val compressionCodec = if (compressionCodecName != null) {
CSVCompressionCodecs.getCodecClassName(compressionCodecName)
} else {
null
}

val maxColumns = 20480

val maxCharsPerColumn = 100000
Expand All @@ -85,7 +96,6 @@ private[sql] case class CSVParameters(@transient parameters: Map[String, String]
}

private[csv] object ParseModes {

val PERMISSIVE_MODE = "PERMISSIVE"
val DROP_MALFORMED_MODE = "DROPMALFORMED"
val FAIL_FAST_MODE = "FAILFAST"
Expand All @@ -107,3 +117,28 @@ private[csv] object ParseModes {
true // We default to permissive is the mode string is not valid
}
}

private[csv] object CSVCompressionCodecs {
val shortCompressionCodecNames = Map(
Copy link
Contributor

@rxin rxin Jan 19, 2016

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

private?

"bzip2" -> classOf[BZip2Codec].getName,
"gzip" -> classOf[GzipCodec].getName,
"lz4" -> classOf[Lz4Codec].getName,
"snappy" -> classOf[SnappyCodec].getName)

/**
* Return the full version of the given codec class.
* If it is already a class name, just return it.
*/
def getCodecClassName(name: String): String = {
val codecName = shortCompressionCodecNames.getOrElse(name.toLowerCase, name)
val codecClassName = try {
// Validate the codec name
Utils.classForName(codecName)
Some(codecName)
} catch {
case e: ClassNotFoundException => None
}
codecClassName.getOrElse(throw new IllegalArgumentException(s"Codec [$codecName] " +
Copy link
Contributor

@aarondav aarondav Jan 20, 2016

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should just put this throw inside the catch block and not bother with the Option stuff

s"is not available. Available codecs are ${shortCompressionCodecNames.keys.mkString(",")}."))
Copy link
Contributor

@rxin rxin Jan 19, 2016

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

add a space after the comma

}
}
10 changes: 10 additions & 0 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/csv/CSVRelation.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ import scala.util.control.NonFatal
import com.google.common.base.Objects
import org.apache.hadoop.fs.{FileStatus, Path}
import org.apache.hadoop.io.{LongWritable, NullWritable, Text}
import org.apache.hadoop.io.SequenceFile.CompressionType
import org.apache.hadoop.mapred.TextInputFormat
import org.apache.hadoop.mapreduce.{Job, TaskAttemptContext}
import org.apache.hadoop.mapreduce.RecordWriter
Expand Down Expand Up @@ -99,6 +100,15 @@ private[csv] class CSVRelation(
}

override def prepareJobForWrite(job: Job): OutputWriterFactory = {
val conf = job.getConfiguration
Option(params.compressionCodec).foreach { codec =>
Copy link
Contributor

@rxin rxin Jan 19, 2016

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

as suggested above, I'd make "compressionCodec" itself an Option, rather than nullable.

conf.set("mapreduce.output.fileoutputformat.compress", "true")
conf.set("mapreduce.output.fileoutputformat.compress.type", CompressionType.BLOCK.toString)
conf.set("mapreduce.output.fileoutputformat.compress.codec", codec)
conf.set("mapreduce.map.output.compress", "true")
conf.set("mapreduce.map.output.compress.codec", codec)
}

new CSVOutputWriterFactory(params)
}

Expand Down
26 changes: 26 additions & 0 deletions sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/csv/CSVSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -349,4 +349,30 @@ class CSVSuite extends QueryTest with SharedSQLContext with SQLTestUtils {
assert(results(0).toSeq === Array(2012, "Tesla", "S", "null", "null"))
assert(results(2).toSeq === Array(null, "Chevy", "Volt", null, null))
}

test("save csv with compression codec option") {
withTempDir { dir =>
val csvDir = new File(dir, "csv").getCanonicalPath
val cars = sqlContext.read
.format("csv")
.option("header", "true")
.load(testFile(carsFile))

cars.coalesce(1).write
.format("csv")
.option("header", "true")
.option("compression", "gZiP")
.save(csvDir)

val compressedFiles = new File(csvDir).listFiles()
assert(compressedFiles.exists(_.getName.endsWith(".gz")))

val carsCopy = sqlContext.read
.format("csv")
.option("header", "true")
.load(csvDir)

verifyCars(carsCopy, withHeader = true)
}
}
}