From d0667e4fc21d35e31f9a48b00ec730f85033e513 Mon Sep 17 00:00:00 2001 From: Mark Cooper Date: Mon, 20 Mar 2017 13:09:00 -0700 Subject: [PATCH 1/4] Fixed lack of support for DATE types. The existing Avro serialization library didn't support DATE typed fields. Folded in additional Databricks Avro lib components and moved the Avro libs to a .cloned package to avoid conflicts. --- build.sbt | 4 +- project/plugins.sbt | 1 + .../databricks/spark/avro/DefaultSource.scala | 289 ++++++++++++++++++ .../spark/avro/clone/AvroOutputWriter.scala | 167 ++++++++++ .../avro/clone/AvroOutputWriterFactory.scala | 36 +++ .../avro/{ => clone}/SchemaConverters.scala | 8 +- .../databricks/spark/avro/clone/package.scala | 36 +++ .../com/spotify/spark/bigquery/package.scala | 3 +- 8 files changed, 538 insertions(+), 6 deletions(-) create mode 100644 src/main/scala/com/databricks/spark/avro/DefaultSource.scala create mode 100644 src/main/scala/com/databricks/spark/avro/clone/AvroOutputWriter.scala create mode 100644 src/main/scala/com/databricks/spark/avro/clone/AvroOutputWriterFactory.scala rename src/main/scala/com/databricks/spark/avro/{ => clone}/SchemaConverters.scala (98%) create mode 100644 src/main/scala/com/databricks/spark/avro/clone/package.scala diff --git a/build.sbt b/build.sbt index 40ae2d2..12ab0ef 100644 --- a/build.sbt +++ b/build.sbt @@ -17,7 +17,7 @@ name := "spark-bigquery" organization := "com.spotify" -scalaVersion := "2.10.6" +scalaVersion := "2.11.8" crossScalaVersions := Seq("2.10.6", "2.11.8") spName := "spotify/spark-bigquery" @@ -27,7 +27,7 @@ spAppendScalaVersion := true spIncludeMaven := true libraryDependencies ++= Seq( - "com.databricks" %% "spark-avro" % "3.0.0", +// "com.databricks" %% "spark-avro" % "3.2.0", "com.google.cloud.bigdataoss" % "bigquery-connector" % "0.7.5-hadoop2" exclude ("com.google.guava", "guava-jdk5"), "org.slf4j" % "slf4j-simple" % "1.7.21", diff --git a/project/plugins.sbt b/project/plugins.sbt index 46d1cb0..c6a83eb 100644 --- a/project/plugins.sbt +++ b/project/plugins.sbt @@ -5,3 +5,4 @@ addSbtPlugin("org.scalastyle" % "scalastyle-sbt-plugin" % "0.8.0") addSbtPlugin("org.scoverage" % "sbt-scoverage" % "1.3.5") addSbtPlugin("org.spark-packages" % "sbt-spark-package" % "0.2.4") addSbtPlugin("org.xerial.sbt" % "sbt-sonatype" % "1.1") +addSbtPlugin("net.virtual-void" % "sbt-dependency-graph" % "0.8.2") diff --git a/src/main/scala/com/databricks/spark/avro/DefaultSource.scala b/src/main/scala/com/databricks/spark/avro/DefaultSource.scala new file mode 100644 index 0000000..0d481e0 --- /dev/null +++ b/src/main/scala/com/databricks/spark/avro/DefaultSource.scala @@ -0,0 +1,289 @@ +/* + * Copyright 2014 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.avro + +import java.io._ +import java.net.URI +import java.util.zip.Deflater + +import com.databricks.spark.avro.DefaultSource.{IgnoreFilesWithoutExtensionProperty, SerializableConfiguration} +import com.databricks.spark.avro.clone.{AvroOutputWriterFactory, SchemaConverters} +import com.esotericsoftware.kryo.io.{Input, Output} +import com.esotericsoftware.kryo.{Kryo, KryoSerializable} +import org.apache.avro.SchemaBuilder +import org.apache.avro.file.{DataFileConstants, DataFileReader} +import org.apache.avro.generic.{GenericDatumReader, GenericRecord} +import org.apache.avro.mapred.{AvroOutputFormat, FsInput} +import org.apache.avro.mapreduce.AvroJob +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.{FileStatus, Path} +import org.apache.hadoop.mapreduce.Job +import org.apache.spark.TaskContext +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.encoders.RowEncoder +import org.apache.spark.sql.catalyst.expressions.GenericRow +import org.apache.spark.sql.execution.datasources.{FileFormat, OutputWriterFactory, PartitionedFile} +import org.apache.spark.sql.sources.{DataSourceRegister, Filter} +import org.apache.spark.sql.types.StructType +import org.slf4j.LoggerFactory + +import scala.util.control.NonFatal + +class DefaultSource extends FileFormat with DataSourceRegister { + private val log = LoggerFactory.getLogger(getClass) + + override def equals(other: Any): Boolean = other match { + case _: DefaultSource => true + case _ => false + } + + override def inferSchema( + spark: SparkSession, + options: Map[String, String], + files: Seq[FileStatus]): Option[StructType] = { + val conf = spark.sparkContext.hadoopConfiguration + + // Schema evolution is not supported yet. Here we only pick a single random sample file to + // figure out the schema of the whole dataset. + val sampleFile = if (conf.getBoolean(IgnoreFilesWithoutExtensionProperty, true)) { + files.find(_.getPath.getName.endsWith(".avro")).getOrElse { + throw new FileNotFoundException( + "No Avro files found. Hadoop option \"avro.mapred.ignore.inputs.without.extension\" is " + + "set to true. Do all input files have \".avro\" extension?" + ) + } + } else { + files.headOption.getOrElse { + throw new FileNotFoundException("No Avro files found.") + } + } + + val avroSchema = { + val in = new FsInput(sampleFile.getPath, conf) + try { + val reader = DataFileReader.openReader(in, new GenericDatumReader[GenericRecord]()) + try { + reader.getSchema + } finally { + reader.close() + } + } finally { + in.close() + } + } + + SchemaConverters.toSqlType(avroSchema).dataType match { + case t: StructType => Some(t) + case _ => throw new RuntimeException( + s"""Avro schema cannot be converted to a Spark SQL StructType: + | + |${avroSchema.toString(true)} + |""".stripMargin) + } + } + + override def shortName(): String = "avro" + + override def prepareWrite( + spark: SparkSession, + job: Job, + options: Map[String, String], + dataSchema: StructType): OutputWriterFactory = { + val recordName = options.getOrElse("recordName", "topLevelRecord") + val recordNamespace = options.getOrElse("recordNamespace", "") + val build = SchemaBuilder.record(recordName).namespace(recordNamespace) + val outputAvroSchema = SchemaConverters.convertStructToAvro(dataSchema, build, recordNamespace) + + AvroJob.setOutputKeySchema(job, outputAvroSchema) + val AVRO_COMPRESSION_CODEC = "spark.sql.avro.compression.codec" + val AVRO_DEFLATE_LEVEL = "spark.sql.avro.deflate.level" + val COMPRESS_KEY = "mapred.output.compress" + + spark.conf.get(AVRO_COMPRESSION_CODEC, "snappy") match { + case "uncompressed" => + log.info("writing uncompressed Avro records") + job.getConfiguration.setBoolean(COMPRESS_KEY, false) + + case "snappy" => + log.info("compressing Avro output using Snappy") + job.getConfiguration.setBoolean(COMPRESS_KEY, true) + job.getConfiguration.set(AvroJob.CONF_OUTPUT_CODEC, DataFileConstants.SNAPPY_CODEC) + + case "deflate" => + val deflateLevel = spark.conf.get( + AVRO_DEFLATE_LEVEL, Deflater.DEFAULT_COMPRESSION.toString).toInt + log.info(s"compressing Avro output using deflate (level=$deflateLevel)") + job.getConfiguration.setBoolean(COMPRESS_KEY, true) + job.getConfiguration.set(AvroJob.CONF_OUTPUT_CODEC, DataFileConstants.DEFLATE_CODEC) + job.getConfiguration.setInt(AvroOutputFormat.DEFLATE_LEVEL_KEY, deflateLevel) + + case unknown: String => + log.error(s"unsupported compression codec $unknown") + } + + new AvroOutputWriterFactory(dataSchema, recordName, recordNamespace) + } + + override def buildReader( + spark: SparkSession, + dataSchema: StructType, + partitionSchema: StructType, + requiredSchema: StructType, + filters: Seq[Filter], + options: Map[String, String], + hadoopConf: Configuration): (PartitionedFile) => Iterator[InternalRow] = { + + val broadcastedConf = + spark.sparkContext.broadcast(new SerializableConfiguration(hadoopConf)) + + (file: PartitionedFile) => { + val log = LoggerFactory.getLogger(classOf[DefaultSource]) + val conf = broadcastedConf.value.value + + // TODO Removes this check once `FileFormat` gets a general file filtering interface method. + // Doing input file filtering is improper because we may generate empty tasks that process no + // input files but stress the scheduler. We should probably add a more general input file + // filtering mechanism for `FileFormat` data sources. See SPARK-16317. + if ( + conf.getBoolean(IgnoreFilesWithoutExtensionProperty, true) && + !file.filePath.endsWith(".avro") + ) { + Iterator.empty + } else { + val reader = { + val in = new FsInput(new Path(new URI(file.filePath)), conf) + try { + DataFileReader.openReader(in, new GenericDatumReader[GenericRecord]()) + } catch { + case NonFatal(e) => + log.error("Exception while opening DataFileReader", e) + in.close() + throw e + } + } + + // Ensure that the reader is closed even if the task fails or doesn't consume the entire + // iterator of records. + Option(TaskContext.get()).foreach { taskContext => + taskContext.addTaskCompletionListener { _ => + reader.close() + } + } + + val fieldExtractors = { + val avroSchema = reader.getSchema + requiredSchema.zipWithIndex.map { case (field, index) => + val avroField = Option(avroSchema.getField(field.name)).getOrElse { + throw new IllegalArgumentException( + s"""Cannot find required column ${field.name} in Avro schema:" + | + |${avroSchema.toString(true)} + """.stripMargin + ) + } + + val converter = SchemaConverters.createConverterToSQL(avroField.schema()) + + (record: GenericRecord, buffer: Array[Any]) => { + buffer(index) = converter(record.get(avroField.pos())) + } + } + } + + new Iterator[InternalRow] { + private val rowBuffer = Array.fill[Any](requiredSchema.length)(null) + + private val safeDataRow = new GenericRow(rowBuffer) + + // Used to convert `Row`s containing data columns into `InternalRow`s. + private val encoderForDataColumns = RowEncoder(requiredSchema) + + private[this] var completed = false + + override def hasNext: Boolean = { + if (completed) { + false + } else { + val r = reader.hasNext + if (!r) { + reader.close() + completed = true + } + r + } + } + + override def next(): InternalRow = { + val record = reader.next() + + var i = 0 + while (i < requiredSchema.length) { + fieldExtractors(i)(record, rowBuffer) + i += 1 + } + + encoderForDataColumns.toRow(safeDataRow) + } + } + } + } + } +} + +object DefaultSource { + val IgnoreFilesWithoutExtensionProperty = "avro.mapred.ignore.inputs.without.extension" + + class SerializableConfiguration(@transient var value: Configuration) + extends Serializable with KryoSerializable { + @transient private[avro] 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)) + } + } +} diff --git a/src/main/scala/com/databricks/spark/avro/clone/AvroOutputWriter.scala b/src/main/scala/com/databricks/spark/avro/clone/AvroOutputWriter.scala new file mode 100644 index 0000000..162b52e --- /dev/null +++ b/src/main/scala/com/databricks/spark/avro/clone/AvroOutputWriter.scala @@ -0,0 +1,167 @@ +/* + * Copyright 2014 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.avro.clone + +import java.io.{IOException, OutputStream} +import java.nio.ByteBuffer +import java.sql.{Date, Timestamp} +import java.util.HashMap + +import org.apache.hadoop.fs.Path + +import scala.collection.immutable.Map +import org.apache.avro.generic.GenericData.Record +import org.apache.avro.generic.GenericRecord +import org.apache.avro.{Schema, SchemaBuilder} +import org.apache.avro.mapred.AvroKey +import org.apache.avro.mapreduce.AvroKeyOutputFormat +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.io.NullWritable +import org.apache.hadoop.mapreduce.{RecordWriter, TaskAttemptContext, TaskAttemptID} +import org.apache.spark.sql.Row +import org.apache.spark.sql.execution.datasources.OutputWriter +import org.apache.spark.sql.types._ + +// NOTE: This class is instantiated and used on executor side only, no need to be serializable. +class AvroOutputWriter( + path: String, + context: TaskAttemptContext, + schema: StructType, + recordName: String, + recordNamespace: String) extends OutputWriter { + + private lazy val converter = createConverterToAvro(schema, recordName, recordNamespace) + + /** + * Overrides the couple of methods responsible for generating the output streams / files so + * that the data can be correctly partitioned + */ + private val recordWriter: RecordWriter[AvroKey[GenericRecord], NullWritable] = + new AvroKeyOutputFormat[GenericRecord]() { + + private def getConfigurationFromContext(context: TaskAttemptContext): Configuration = { + // Use reflection to get the Configuration. This is necessary because TaskAttemptContext + // is a class in Hadoop 1.x and an interface in Hadoop 2.x. + val method = context.getClass.getMethod("getConfiguration") + method.invoke(context).asInstanceOf[Configuration] + } + + override def getDefaultWorkFile(context: TaskAttemptContext, extension: String): Path = { + val uniqueWriteJobId = + getConfigurationFromContext(context).get("spark.sql.sources.writeJobUUID") + val taskAttemptId: TaskAttemptID = { + // Use reflection to get the TaskAttemptID. This is necessary because TaskAttemptContext + // is a class in Hadoop 1.x and an interface in Hadoop 2.x. + val method = context.getClass.getMethod("getTaskAttemptID") + method.invoke(context).asInstanceOf[TaskAttemptID] + } + val split = taskAttemptId.getTaskID.getId + new Path(path, f"part-r-$split%05d-$uniqueWriteJobId$extension") + } + + @throws(classOf[IOException]) + override def getAvroFileOutputStream(c: TaskAttemptContext): OutputStream = { + val path = getDefaultWorkFile(context, ".avro") + path.getFileSystem(getConfigurationFromContext(context)).create(path) + } + + }.getRecordWriter(context) + + override def write(row: Row): Unit = { + val key = new AvroKey(converter(row).asInstanceOf[GenericRecord]) + recordWriter.write(key, NullWritable.get()) + } + + override def close(): Unit = recordWriter.close(context) + + // scalastyle:off + + /** + * This function constructs converter function for a given sparkSQL datatype. This is used in + * writing Avro records out to disk + */ + private def createConverterToAvro( + dataType: DataType, + structName: String, + recordNamespace: String): (Any) => Any = { + dataType match { + case BinaryType => (item: Any) => item match { + case null => null + case bytes: Array[Byte] => ByteBuffer.wrap(bytes) + } + case ByteType | ShortType | IntegerType | LongType | + FloatType | DoubleType | StringType | BooleanType => identity + case _: DecimalType => (item: Any) => if (item == null) null else item.toString + case TimestampType => (item: Any) => + if (item == null) null else item.asInstanceOf[Timestamp].getTime + case DateType => (item: Any) => + if (item == null) null else item.asInstanceOf[Date].getTime + case ArrayType(elementType, _) => + val elementConverter = createConverterToAvro(elementType, structName, recordNamespace) + (item: Any) => { + if (item == null) { + null + } else { + val sourceArray = item.asInstanceOf[Seq[Any]] + val sourceArraySize = sourceArray.size + val targetArray = new Array[Any](sourceArraySize) + var idx = 0 + while (idx < sourceArraySize) { + targetArray(idx) = elementConverter(sourceArray(idx)) + idx += 1 + } + targetArray + } + } + case MapType(StringType, valueType, _) => + val valueConverter = createConverterToAvro(valueType, structName, recordNamespace) + (item: Any) => { + if (item == null) { + null + } else { + val javaMap = new HashMap[String, Any]() + item.asInstanceOf[Map[String, Any]].foreach { case (key, value) => + javaMap.put(key, valueConverter(value)) + } + javaMap + } + } + case structType: StructType => + val builder = SchemaBuilder.record(structName).namespace(recordNamespace) + val schema: Schema = SchemaConverters.convertStructToAvro( + structType, builder, recordNamespace) + val fieldConverters = structType.fields.map(field => + createConverterToAvro(field.dataType, field.name, recordNamespace)) + (item: Any) => { + if (item == null) { + null + } else { + val record = new Record(schema) + val convertersIterator = fieldConverters.iterator + val fieldNamesIterator = dataType.asInstanceOf[StructType].fieldNames.iterator + val rowIterator = item.asInstanceOf[Row].toSeq.iterator + + while (convertersIterator.hasNext) { + val converter = convertersIterator.next() + record.put(fieldNamesIterator.next(), converter(rowIterator.next())) + } + record + } + } + } + } +} diff --git a/src/main/scala/com/databricks/spark/avro/clone/AvroOutputWriterFactory.scala b/src/main/scala/com/databricks/spark/avro/clone/AvroOutputWriterFactory.scala new file mode 100644 index 0000000..07313bd --- /dev/null +++ b/src/main/scala/com/databricks/spark/avro/clone/AvroOutputWriterFactory.scala @@ -0,0 +1,36 @@ +/* + * Copyright 2014 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.avro.clone + +import org.apache.hadoop.mapreduce.TaskAttemptContext + +import org.apache.spark.sql.execution.datasources.{OutputWriter, OutputWriterFactory} +import org.apache.spark.sql.types.StructType + +class AvroOutputWriterFactory( + schema: StructType, + recordName: String, + recordNamespace: String) extends OutputWriterFactory { + + def newInstance( + path: String, + bucketId: Option[Int], + dataSchema: StructType, + context: TaskAttemptContext): OutputWriter = { + new AvroOutputWriter(path, context, schema, recordName, recordNamespace) + } +} diff --git a/src/main/scala/com/databricks/spark/avro/SchemaConverters.scala b/src/main/scala/com/databricks/spark/avro/clone/SchemaConverters.scala similarity index 98% rename from src/main/scala/com/databricks/spark/avro/SchemaConverters.scala rename to src/main/scala/com/databricks/spark/avro/clone/SchemaConverters.scala index 350d117..b5f2011 100644 --- a/src/main/scala/com/databricks/spark/avro/SchemaConverters.scala +++ b/src/main/scala/com/databricks/spark/avro/clone/SchemaConverters.scala @@ -13,8 +13,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -// scalastyle:off -package com.databricks.spark.avro +package com.databricks.spark.avro.clone import java.nio.ByteBuffer import java.util.HashMap @@ -29,6 +28,8 @@ import org.apache.avro.Schema.Type._ import org.apache.spark.sql.Row import org.apache.spark.sql.types._ +// scalastyle:off + /** * This object contains method that are used to convert sparkSQL schemas to avro schemas and vice * versa. @@ -225,6 +226,7 @@ object SchemaConverters { case BinaryType => schemaBuilder.bytesType() case BooleanType => schemaBuilder.booleanType() case TimestampType => schemaBuilder.longType() + case DateType => schemaBuilder.longType() case ArrayType(elementType, _) => val builder = getSchemaBuilder(dataType.asInstanceOf[ArrayType].containsNull) @@ -268,6 +270,7 @@ object SchemaConverters { case BinaryType => newFieldBuilder.bytesType() case BooleanType => newFieldBuilder.booleanType() case TimestampType => newFieldBuilder.longType() + case DateType => newFieldBuilder.longType() case ArrayType(elementType, _) => val builder = getSchemaBuilder(dataType.asInstanceOf[ArrayType].containsNull) @@ -297,4 +300,3 @@ object SchemaConverters { } } } -// scalastyle:on \ No newline at end of file diff --git a/src/main/scala/com/databricks/spark/avro/clone/package.scala b/src/main/scala/com/databricks/spark/avro/clone/package.scala new file mode 100644 index 0000000..c249610 --- /dev/null +++ b/src/main/scala/com/databricks/spark/avro/clone/package.scala @@ -0,0 +1,36 @@ +/* + * Copyright 2014 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.avro + +import org.apache.spark.sql.{DataFrame, DataFrameReader, DataFrameWriter} + +package object clone { + /** + * Adds a method, `avro`, to DataFrameWriter that allows you to write avro files using + * the DataFileWriter + */ + implicit class AvroDataFrameWriter[T](writer: DataFrameWriter[T]) { + def avro: String => Unit = writer.format("com.databricks.spark.avro").save + } + + /** + * Adds a method, `avro`, to DataFrameReader that allows you to read avro files using + * the DataFileReade + */ + implicit class AvroDataFrameReader(reader: DataFrameReader) { + def avro: String => DataFrame = reader.format("com.databricks.spark.avro").load + } +} diff --git a/src/main/scala/com/spotify/spark/bigquery/package.scala b/src/main/scala/com/spotify/spark/bigquery/package.scala index d295c4e..5dda829 100644 --- a/src/main/scala/com/spotify/spark/bigquery/package.scala +++ b/src/main/scala/com/spotify/spark/bigquery/package.scala @@ -17,7 +17,8 @@ package com.spotify.spark -import com.databricks.spark.avro._ +import com.databricks.spark.avro.clone._ +import com.databricks.spark.avro.clone.SchemaConverters import com.google.api.services.bigquery.model.TableReference import com.google.cloud.hadoop.fs.gcs.GoogleHadoopFileSystem import com.google.cloud.hadoop.io.bigquery._ From b4ecc255519dcd81a99433b8af1583a1b14ce9e8 Mon Sep 17 00:00:00 2001 From: Mark Cooper Date: Mon, 20 Mar 2017 13:32:16 -0700 Subject: [PATCH 2/4] Bump version to: 0.2.1-DATEFIX --- version.sbt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/version.sbt b/version.sbt index cfeb331..8fd77af 100644 --- a/version.sbt +++ b/version.sbt @@ -1 +1 @@ -version in ThisBuild := "0.2.1-SNAPSHOT" +version in ThisBuild := "0.2.1-DATEFIX" From 4ae08624d77ff75141cde58e8d978734bfb6ceb8 Mon Sep 17 00:00:00 2001 From: Mark Cooper Date: Wed, 26 Apr 2017 17:22:15 -0700 Subject: [PATCH 3/4] Support for BigQuery Timestamps with Avro encoded data Horible, ugly not so good hacks * Convert DateType => String (this type isn't often used, we use it for birth dates) * Encode TimestampType in microsecond * Generate and pass a long a schema to BigQuery * Added a handy bigQueryTableExists() util method --- scalastyle-config.xml | 2 +- .../spark/avro/clone/AvroOutputWriter.scala | 4 +- .../spark/avro/clone/SchemaConverters.scala | 4 +- .../spark/bigquery/BigQueryClient.scala | 11 ++- .../com/spotify/spark/bigquery/package.scala | 22 +++-- .../scala/spark/bigquery/BigQuerySchema.scala | 80 +++++++++++++++++++ version.sbt | 2 +- 7 files changed, 113 insertions(+), 12 deletions(-) create mode 100644 src/main/scala/spark/bigquery/BigQuerySchema.scala diff --git a/scalastyle-config.xml b/scalastyle-config.xml index 2b10c34..1806655 100644 --- a/scalastyle-config.xml +++ b/scalastyle-config.xml @@ -32,7 +32,7 @@ - + diff --git a/src/main/scala/com/databricks/spark/avro/clone/AvroOutputWriter.scala b/src/main/scala/com/databricks/spark/avro/clone/AvroOutputWriter.scala index 162b52e..698c929 100644 --- a/src/main/scala/com/databricks/spark/avro/clone/AvroOutputWriter.scala +++ b/src/main/scala/com/databricks/spark/avro/clone/AvroOutputWriter.scala @@ -107,9 +107,9 @@ class AvroOutputWriter( FloatType | DoubleType | StringType | BooleanType => identity case _: DecimalType => (item: Any) => if (item == null) null else item.toString case TimestampType => (item: Any) => - if (item == null) null else item.asInstanceOf[Timestamp].getTime + if (item == null) null else item.asInstanceOf[Timestamp].getTime * 1000l // MNC: mills -> micros hack. case DateType => (item: Any) => - if (item == null) null else item.asInstanceOf[Date].getTime + if (item == null) null else item.asInstanceOf[Date].toString case ArrayType(elementType, _) => val elementConverter = createConverterToAvro(elementType, structName, recordNamespace) (item: Any) => { diff --git a/src/main/scala/com/databricks/spark/avro/clone/SchemaConverters.scala b/src/main/scala/com/databricks/spark/avro/clone/SchemaConverters.scala index b5f2011..45bef66 100644 --- a/src/main/scala/com/databricks/spark/avro/clone/SchemaConverters.scala +++ b/src/main/scala/com/databricks/spark/avro/clone/SchemaConverters.scala @@ -226,7 +226,7 @@ object SchemaConverters { case BinaryType => schemaBuilder.bytesType() case BooleanType => schemaBuilder.booleanType() case TimestampType => schemaBuilder.longType() - case DateType => schemaBuilder.longType() + case DateType => schemaBuilder.stringType() case ArrayType(elementType, _) => val builder = getSchemaBuilder(dataType.asInstanceOf[ArrayType].containsNull) @@ -270,7 +270,7 @@ object SchemaConverters { case BinaryType => newFieldBuilder.bytesType() case BooleanType => newFieldBuilder.booleanType() case TimestampType => newFieldBuilder.longType() - case DateType => newFieldBuilder.longType() + case DateType => newFieldBuilder.stringType() case ArrayType(elementType, _) => val builder = getSchemaBuilder(dataType.asInstanceOf[ArrayType].containsNull) diff --git a/src/main/scala/com/spotify/spark/bigquery/BigQueryClient.scala b/src/main/scala/com/spotify/spark/bigquery/BigQueryClient.scala index e14120e..c993b7c 100644 --- a/src/main/scala/com/spotify/spark/bigquery/BigQueryClient.scala +++ b/src/main/scala/com/spotify/spark/bigquery/BigQueryClient.scala @@ -104,18 +104,27 @@ private[bigquery] class BigQueryClient(conf: Configuration) { */ def query(sqlQuery: String): TableReference = queryCache.get(sqlQuery) + /** + * Check if the table exists - useful for pre-write logic. + */ + def tableExists(table: TableReference): Boolean = { + new BigQueryHelper(bigquery).tableExists(table) + } + /** * Load an Avro data set on GCS to a BigQuery table. */ def load(gcsPath: String, destinationTable: TableReference, writeDisposition: WriteDisposition.Value = null, - createDisposition: CreateDisposition.Value = null): Unit = { + createDisposition: CreateDisposition.Value = null, + schema: TableSchema = null): Unit = { val tableName = BigQueryStrings.toString(destinationTable) logger.info(s"Loading $gcsPath into $tableName") var loadConfig = new JobConfigurationLoad() .setDestinationTable(destinationTable) .setSourceFormat("AVRO") .setSourceUris(List(gcsPath + "/*.avro").asJava) + .setSchema(schema) if (writeDisposition != null) { loadConfig = loadConfig.setWriteDisposition(writeDisposition.toString) } diff --git a/src/main/scala/com/spotify/spark/bigquery/package.scala b/src/main/scala/com/spotify/spark/bigquery/package.scala index 5dda829..44848fc 100644 --- a/src/main/scala/com/spotify/spark/bigquery/package.scala +++ b/src/main/scala/com/spotify/spark/bigquery/package.scala @@ -19,7 +19,7 @@ package com.spotify.spark import com.databricks.spark.avro.clone._ import com.databricks.spark.avro.clone.SchemaConverters -import com.google.api.services.bigquery.model.TableReference +import com.google.api.services.bigquery.model.{TableReference, TableSchema} import com.google.cloud.hadoop.fs.gcs.GoogleHadoopFileSystem import com.google.cloud.hadoop.io.bigquery._ import org.apache.avro.Schema @@ -148,17 +148,27 @@ package object bigquery { sqlContext.setConf("spark.sql.avro.compression.codec", "deflate") + /** + * Check for existence of a table. + * @param tableRef + * @return + */ + def bigQueryTableExists(tableRef: TableReference): Boolean = { + bq.tableExists(tableRef) + } + /** * Save a [[DataFrame]] to a BigQuery table. */ def saveAsBigQueryTable(tableRef: TableReference, writeDisposition: WriteDisposition.Value, - createDisposition: CreateDisposition.Value): Unit = { + createDisposition: CreateDisposition.Value, + schema: TableSchema): Unit = { val bucket = conf.get(BigQueryConfiguration.GCS_BUCKET_KEY) val temp = s"spark-bigquery-${System.currentTimeMillis()}=${Random.nextInt(Int.MaxValue)}" val gcsPath = s"gs://$bucket/hadoop/tmp/spark-bigquery/$temp" self.write.avro(gcsPath) - val df = bq.load(gcsPath, tableRef, writeDisposition, createDisposition) + val df = bq.load(gcsPath, tableRef, writeDisposition, createDisposition, schema) delete(new Path(gcsPath)) df } @@ -168,11 +178,13 @@ package object bigquery { */ def saveAsBigQueryTable(tableSpec: String, writeDisposition: WriteDisposition.Value = null, - createDisposition: CreateDisposition.Value = null): Unit = + createDisposition: CreateDisposition.Value = null, + schema: TableSchema = null): Unit = saveAsBigQueryTable( BigQueryStrings.parseTableReference(tableSpec), writeDisposition, - createDisposition) + createDisposition, + schema) private def delete(path: Path): Unit = { val fs = FileSystem.get(path.toUri, conf) diff --git a/src/main/scala/spark/bigquery/BigQuerySchema.scala b/src/main/scala/spark/bigquery/BigQuerySchema.scala new file mode 100644 index 0000000..9a03dcf --- /dev/null +++ b/src/main/scala/spark/bigquery/BigQuerySchema.scala @@ -0,0 +1,80 @@ +/* + * Copyright 2016 Appsflyer. + * + * 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.appsflyer.spark.bigquery + +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.types._ +import org.json4s.JsonAST.{JArray, JValue} +import org.json4s.JsonDSL._ +import org.json4s.jackson.JsonMethods.{pretty, render} + +/** + * Builds BigQuery input JSON schema based on DataFrame. + * Example schema can be found here: https://cloud.google.com/bigquery/docs/personsDataSchema.json + */ +object BigQuerySchema { + + private def getMode(field: StructField) = { + field.dataType match { + case ArrayType(_, _) => "REPEATED" + case _ => if (field.nullable) "NULLABLE" else "REQUIRED" + } + } + + private def getTypeName(dataType: DataType) = { + dataType match { + case ByteType | ShortType | IntegerType | LongType => "INTEGER" + case FloatType | DoubleType => "FLOAT" + case _: DecimalType | StringType => "STRING" + case BinaryType => "BYTES" + case BooleanType => "BOOLEAN" + case TimestampType => "TIMESTAMP" + case ArrayType(_, _) | MapType(_, _, _) | _: StructType => "RECORD" + case DateType => "STRING" + } + } + + private def typeToJson(field: StructField, dataType: DataType): JValue = { + dataType match { + case structType: StructType => + ("type" -> getTypeName(dataType)) ~ + ("fields" -> structType.fields.map(fieldToJson(_)).toList) + case arrayType: ArrayType => + arrayType.elementType match { + case _: ArrayType => + throw new IllegalArgumentException(s"Multidimensional arrays are not supported: ${field.name}") + case other => + typeToJson(field, other) + } + case mapType: MapType => + throw new IllegalArgumentException(s"Unsupported type: ${dataType}") + case other => + ("type" -> getTypeName(dataType)) + } + } + + private def fieldToJson(field: StructField): JValue = { + ("name" -> field.name) ~ + ("mode" -> getMode(field)) merge + typeToJson(field, field.dataType) + } + + def apply(df: DataFrame): String = { + pretty(render(JArray(df.schema.fields.map(fieldToJson(_)).toList))) + } +} diff --git a/version.sbt b/version.sbt index 8fd77af..74410b2 100644 --- a/version.sbt +++ b/version.sbt @@ -1 +1 @@ -version in ThisBuild := "0.2.1-DATEFIX" +version in ThisBuild := "0.2.1-BRIGADE02" From 8cfd1af89e216b01b80b019245f3450e6ba4e019 Mon Sep 17 00:00:00 2001 From: Mark Cooper Date: Thu, 27 Apr 2017 13:04:03 -0700 Subject: [PATCH 4/4] Adds a new saveAsBigQueryTableWithRichSchema method This method extracts the schema from the Dataframe and uses it when reading in and generating the table schema. The old mechanism simply uses the Avro defition to infer the table schema. --- .../com/spotify/spark/bigquery/package.scala | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/src/main/scala/com/spotify/spark/bigquery/package.scala b/src/main/scala/com/spotify/spark/bigquery/package.scala index 44848fc..a5d34c0 100644 --- a/src/main/scala/com/spotify/spark/bigquery/package.scala +++ b/src/main/scala/com/spotify/spark/bigquery/package.scala @@ -17,6 +17,7 @@ package com.spotify.spark +import com.appsflyer.spark.bigquery.BigQuerySchema import com.databricks.spark.avro.clone._ import com.databricks.spark.avro.clone.SchemaConverters import com.google.api.services.bigquery.model.{TableReference, TableSchema} @@ -186,6 +187,26 @@ package object bigquery { createDisposition, schema) + /** + * Save a [[DataFrame]] to a BigQuery table. Extract the schema definition from + * the Dataframe and force BigQuery to use this schema when importing the Avro data. + * This has the benefit that Timestamp and Date fields don't get coerced to longs but + * instead use the types available in BigQuery. + */ + def saveAsBigQueryTableWithRichSchema(tableSpec: String, + writeDisposition: WriteDisposition.Value = null, + createDisposition: CreateDisposition.Value = null): Unit = { + + val schemaJson = BigQuerySchema(self) + val tableSchema = new TableSchema().setFields(BigQueryUtils.getSchemaFromString(schemaJson)) + + saveAsBigQueryTable( + BigQueryStrings.parseTableReference(tableSpec), + writeDisposition, + createDisposition, + tableSchema) + } + private def delete(path: Path): Unit = { val fs = FileSystem.get(path.toUri, conf) fs.delete(path, true)