-
Notifications
You must be signed in to change notification settings - Fork 349
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Use FileFormat-based data source instead of HadoopRDD for reads
This patch refactors this library's read path to use a Spark 2.0's `FileFormat`-based data source to read unloaded Redshift output from S3. This approach has a few advantages over using our existing `HadoopRDD`-based approach: - It will benefit from performance improvements in `FileScanRDD` and `HadoopFsRelation`, including automatic coalescing. - We don't have to create a separate RDD per partition and union them together, making the RDD DAG smaller. The bulk of the diff are helper classes copied from Spark and `spark-avro` and inlined here for API compatibility / stability purposes. Some of the new classes implemented here are likely to become incompatible with new releases of Spark, but note that `spark-avro` itself relies on similar unstable / experimental APIs and thus this library is already vulnerable to changes to those APIs (in other words, this change is not making our compatibility story significantly worse). Author: Josh Rosen <[email protected]> Author: Josh Rosen <[email protected]> Closes #289 from JoshRosen/use-fileformat-for-reads.
- Loading branch information
Showing
11 changed files
with
405 additions
and
70 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
70 changes: 70 additions & 0 deletions
70
src/main/scala/com/databricks/spark/redshift/RecordReaderIterator.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,70 @@ | ||
| /* | ||
| * Licensed to the Apache Software Foundation (ASF) under one or more | ||
| * contributor license agreements. See the NOTICE file distributed with | ||
| * this work for additional information regarding copyright ownership. | ||
| * The ASF licenses this file to You 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 com.databricks.spark.redshift | ||
|
|
||
| import java.io.Closeable | ||
|
|
||
| import org.apache.hadoop.mapreduce.RecordReader | ||
|
|
||
| /** | ||
| * An adaptor from a Hadoop [[RecordReader]] to an [[Iterator]] over the values returned. | ||
| * | ||
| * This is copied from Apache Spark and is inlined here to avoid depending on Spark internals | ||
| * in this external library. | ||
| */ | ||
| private[redshift] class RecordReaderIterator[T]( | ||
| private[this] var rowReader: RecordReader[_, T]) extends Iterator[T] with Closeable { | ||
| private[this] var havePair = false | ||
| private[this] var finished = false | ||
|
|
||
| override def hasNext: Boolean = { | ||
| if (!finished && !havePair) { | ||
| finished = !rowReader.nextKeyValue | ||
| if (finished) { | ||
| // Close and release the reader here; close() will also be called when the task | ||
| // completes, but for tasks that read from many files, it helps to release the | ||
| // resources early. | ||
| close() | ||
| } | ||
| havePair = !finished | ||
| } | ||
| !finished | ||
| } | ||
|
|
||
| override def next(): T = { | ||
| if (!hasNext) { | ||
| throw new java.util.NoSuchElementException("End of stream") | ||
| } | ||
| havePair = false | ||
| rowReader.getCurrentValue | ||
| } | ||
|
|
||
| override def close(): Unit = { | ||
| if (rowReader != null) { | ||
| try { | ||
| // Close the reader and release it. Note: it's very important that we don't close the | ||
| // reader more than once, since that exposes us to MAPREDUCE-5918 when running against | ||
| // older Hadoop 2.x releases. That bug can lead to non-deterministic corruption issues | ||
| // when reading compressed input. | ||
| rowReader.close() | ||
| } finally { | ||
| rowReader = null | ||
| } | ||
| } | ||
| } | ||
| } |
102 changes: 102 additions & 0 deletions
102
src/main/scala/com/databricks/spark/redshift/RedshiftFileFormat.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,102 @@ | ||
| /* | ||
| * Copyright 2016 Databricks | ||
| * | ||
| * 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 com.databricks.spark.redshift | ||
|
|
||
| import java.net.URI | ||
|
|
||
| import org.apache.hadoop.conf.Configuration | ||
| import org.apache.hadoop.fs.{FileStatus, Path} | ||
| import org.apache.hadoop.mapreduce._ | ||
| import org.apache.hadoop.mapreduce.lib.input.FileSplit | ||
| import org.apache.hadoop.mapreduce.task.TaskAttemptContextImpl | ||
| import org.apache.spark.TaskContext | ||
| import org.apache.spark.sql.SparkSession | ||
| import org.apache.spark.sql.catalyst.InternalRow | ||
| import org.apache.spark.sql.execution.datasources._ | ||
| import org.apache.spark.sql.sources.Filter | ||
| import org.apache.spark.sql.types.StructType | ||
|
|
||
| /** | ||
| * Internal data source used for reading Redshift UNLOAD files. | ||
| * | ||
| * This is not intended for public consumption / use outside of this package and therefore | ||
| * no API stability is guaranteed. | ||
| */ | ||
| private[redshift] class RedshiftFileFormat extends FileFormat { | ||
| override def inferSchema( | ||
| sparkSession: SparkSession, | ||
| options: Map[String, String], | ||
| files: Seq[FileStatus]): Option[StructType] = { | ||
| // Schema is provided by caller. | ||
| None | ||
| } | ||
|
|
||
| override def prepareWrite( | ||
| sparkSession: SparkSession, | ||
| job: Job, | ||
| options: Map[String, String], | ||
| dataSchema: StructType): OutputWriterFactory = { | ||
| throw new UnsupportedOperationException(s"prepareWrite is not supported for $this") | ||
| } | ||
|
|
||
| override def isSplitable( | ||
| sparkSession: SparkSession, | ||
| options: Map[String, String], | ||
| path: Path): Boolean = { | ||
| // Our custom InputFormat handles split records properly | ||
| true | ||
| } | ||
|
|
||
| override def buildReader( | ||
| sparkSession: SparkSession, | ||
| dataSchema: StructType, | ||
| partitionSchema: StructType, | ||
| requiredSchema: StructType, | ||
| filters: Seq[Filter], | ||
| options: Map[String, String], | ||
| hadoopConf: Configuration): (PartitionedFile) => Iterator[InternalRow] = { | ||
|
|
||
| require(partitionSchema.isEmpty) | ||
| require(filters.isEmpty) | ||
| require(dataSchema == requiredSchema) | ||
|
|
||
| val broadcastedConf = | ||
| sparkSession.sparkContext.broadcast(new SerializableConfiguration(hadoopConf)) | ||
|
|
||
| (file: PartitionedFile) => { | ||
| val conf = broadcastedConf.value.value | ||
|
|
||
| val fileSplit = new FileSplit( | ||
| new Path(new URI(file.filePath)), | ||
| file.start, | ||
| file.length, | ||
| // TODO: Implement Locality | ||
| Array.empty) | ||
| val attemptId = new TaskAttemptID(new TaskID(new JobID(), TaskType.MAP, 0), 0) | ||
| val hadoopAttemptContext = new TaskAttemptContextImpl(conf, attemptId) | ||
| val reader = new RedshiftRecordReader | ||
| reader.initialize(fileSplit, hadoopAttemptContext) | ||
| val iter = new RecordReaderIterator[Array[String]](reader) | ||
| // Ensure that the record reader is closed upon task completion. It will ordinarily | ||
| // be closed once it is completely iterated, but this is necessary to guard against | ||
| // resource leaks in case the task fails or is interrupted. | ||
| Option(TaskContext.get()).foreach(_.addTaskCompletionListener(_ => iter.close())) | ||
| val converter = Conversions.createRowConverter(requiredSchema) | ||
| iter.map(converter) | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
65 changes: 65 additions & 0 deletions
65
src/main/scala/com/databricks/spark/redshift/SerializableConfiguration.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,65 @@ | ||
| /* | ||
| * Copyright 2016 Databricks | ||
| * | ||
| * 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 com.databricks.spark.redshift | ||
|
|
||
| import java.io._ | ||
|
|
||
| import com.esotericsoftware.kryo.io.{Input, Output} | ||
| import com.esotericsoftware.kryo.{Kryo, KryoSerializable} | ||
| import org.apache.hadoop.conf.Configuration | ||
| import org.slf4j.LoggerFactory | ||
|
|
||
| import scala.util.control.NonFatal | ||
|
|
||
| class SerializableConfiguration(@transient var value: Configuration) | ||
| extends Serializable with KryoSerializable { | ||
| @transient private[redshift] lazy val log = LoggerFactory.getLogger(getClass) | ||
|
|
||
| private def writeObject(out: ObjectOutputStream): Unit = tryOrIOException { | ||
| out.defaultWriteObject() | ||
| value.write(out) | ||
| } | ||
|
|
||
| private def readObject(in: ObjectInputStream): Unit = tryOrIOException { | ||
| value = new Configuration(false) | ||
| value.readFields(in) | ||
| } | ||
|
|
||
| private def tryOrIOException[T](block: => T): T = { | ||
| try { | ||
| block | ||
| } catch { | ||
| case e: IOException => | ||
| log.error("Exception encountered", e) | ||
| throw e | ||
| case NonFatal(e) => | ||
| log.error("Exception encountered", e) | ||
| throw new IOException(e) | ||
| } | ||
| } | ||
|
|
||
| def write(kryo: Kryo, out: Output): Unit = { | ||
| val dos = new DataOutputStream(out) | ||
| value.write(dos) | ||
| dos.flush() | ||
| } | ||
|
|
||
| def read(kryo: Kryo, in: Input): Unit = { | ||
| value = new Configuration(false) | ||
| value.readFields(new DataInputStream(in)) | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.