-
Notifications
You must be signed in to change notification settings - Fork 28.2k
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.
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
[SPARK-746][CORE] Added Avro Serialization to Kryo #7004
Changes from 12 commits
97fba62
2b545cc
f4ae251
ab46d10
d421bf5
6d1925c
0f5471a
c5fe794
fa9298b
1183a48
dd71efe
c0cf329
8158d51
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -18,11 +18,12 @@ | |
package org.apache.spark | ||
|
||
import java.util.concurrent.ConcurrentHashMap | ||
import java.util.concurrent.atomic.AtomicBoolean | ||
|
||
import scala.collection.JavaConverters._ | ||
import scala.collection.mutable.LinkedHashSet | ||
|
||
import org.apache.avro.{SchemaNormalization, Schema} | ||
|
||
import org.apache.spark.serializer.KryoSerializer | ||
import org.apache.spark.util.Utils | ||
|
||
|
@@ -161,6 +162,25 @@ class SparkConf(loadDefaults: Boolean) extends Cloneable with Logging { | |
this | ||
} | ||
|
||
private final val avroNamespace = "avro.schema." | ||
|
||
/** | ||
* Use Kryo serialization and register the given set of Avro schemas so that the generic | ||
* record serializer can decrease network IO | ||
*/ | ||
def registerAvroSchemas(schemas: Schema*): SparkConf = { | ||
schemas.foldLeft(this) { (conf, schema) => | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Why do you need a fold here? This seems confusing. Why not just loop over the schemas and call There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Got it, changed. |
||
conf.set(avroNamespace + SchemaNormalization.parsingFingerprint64(schema), schema.toString) | ||
} | ||
} | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. can you start the body of the method on the next line, and wrap the entire method w/ braces? |
||
|
||
/** Gets all the avro schemas in the configuration used in the generic Avro record serializer */ | ||
def getAvroSchema: Map[Long, String] = { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. What do the keys and values of this map denote? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The keys are longs, which represent the a unique ID of the schema and the values are the string representation of the schema. |
||
getAll.filter { case (k, v) => k.startsWith(avroNamespace) } | ||
.map { case (k, v) => (k.substring(avroNamespace.length).toLong, v) } | ||
.toMap | ||
} | ||
|
||
/** Remove a parameter from the configuration */ | ||
def remove(key: String): SparkConf = { | ||
settings.remove(key) | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,140 @@ | ||
/* | ||
* 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 org.apache.spark.serializer | ||
|
||
import java.io.{ByteArrayInputStream, ByteArrayOutputStream} | ||
import java.nio.ByteBuffer | ||
|
||
import scala.collection.mutable | ||
|
||
import com.esotericsoftware.kryo.{Kryo, Serializer => KSerializer} | ||
import com.esotericsoftware.kryo.io.{Input => KryoInput, Output => KryoOutput} | ||
import org.apache.avro.{Schema, SchemaNormalization} | ||
import org.apache.avro.generic.{GenericData, GenericRecord} | ||
import org.apache.avro.io._ | ||
import org.apache.commons.io.IOUtils | ||
|
||
import org.apache.spark.SparkEnv | ||
import org.apache.spark.io.CompressionCodec | ||
|
||
/** | ||
* Custom serializer used for generic Avro records. If the user registers the schemas | ||
* ahead of time, then the schema's fingerprint will be sent with each message instead of the actual | ||
* schema, as to reduce network IO. | ||
* Actions like parsing or compressing schemas are computationally expensive so the serializer | ||
* caches all previously seen values as to reduce the amount of work needed to do. | ||
*/ | ||
private[serializer] class GenericAvroSerializer(schemas: Map[Long, String]) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Similar comment here: what are the types in There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. See above, but I have also updated the comments. |
||
extends KSerializer[GenericRecord] { | ||
|
||
/** Used to reduce the amount of effort to compress the schema */ | ||
private val compressCache = new mutable.HashMap[Schema, Array[Byte]]() | ||
private val decompressCache = new mutable.HashMap[ByteBuffer, Schema]() | ||
|
||
/** Reuses the same datum reader/writer since the same schema will be used many times */ | ||
private val writerCache = new mutable.HashMap[Schema, DatumWriter[_]]() | ||
private val readerCache = new mutable.HashMap[Schema, DatumReader[_]]() | ||
|
||
/** Fingerprinting is very expensive so this alleviates most of the work */ | ||
private val fingerprintCache = new mutable.HashMap[Schema, Long]() | ||
private val schemaCache = new mutable.HashMap[Long, Schema]() | ||
|
||
/** This needs to be lazy since SparkEnv is not initialized yet sometimes when this is called */ | ||
private lazy val codec = CompressionCodec.createCodec(SparkEnv.get.conf) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. just curious, where does this constructed when there isn't a SparkEnv? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Why not just accept a There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Several of the tests involving block replication fail when this value is not lazily defined. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I originally had it accept SparkConf but I was getting serialization errors since SparkConf is not serializable. This way prevents having to send the configuration with the serializer. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Joe and I talked about this a bit offline -- the reason for this is that There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. for anybody that is curious -- I had gotten myself pretty confused about why this change would make the In practice this means the |
||
|
||
/** | ||
* Used to compress Schemas when they are being sent over the wire. | ||
* The compression results are memoized to reduce the compression time since the | ||
* same schema is compressed many times over | ||
*/ | ||
def compress(schema: Schema): Array[Byte] = compressCache.getOrElseUpdate(schema, { | ||
val bos = new ByteArrayOutputStream() | ||
val out = codec.compressedOutputStream(bos) | ||
out.write(schema.toString.getBytes("UTF-8")) | ||
out.close() | ||
bos.toByteArray | ||
}) | ||
|
||
/** | ||
* Decompresses the schema into the actual in-memory object. Keeps an internal cache of already | ||
* seen values so to limit the number of times that decompression has to be done. | ||
*/ | ||
def decompress(schemaBytes: ByteBuffer): Schema = decompressCache.getOrElseUpdate(schemaBytes, { | ||
val bis = new ByteArrayInputStream(schemaBytes.array()) | ||
val bytes = IOUtils.toByteArray(codec.compressedInputStream(bis)) | ||
new Schema.Parser().parse(new String(bytes, "UTF-8")) | ||
}) | ||
|
||
/** | ||
* Serializes a record to the given output stream. It caches a lot of the internal data as | ||
* to not redo work | ||
*/ | ||
def serializeDatum[R <: GenericRecord](datum: R, output: KryoOutput): Unit = { | ||
val encoder = EncoderFactory.get.binaryEncoder(output, null) | ||
val schema = datum.getSchema | ||
val fingerprint = fingerprintCache.getOrElseUpdate(schema, { | ||
SchemaNormalization.parsingFingerprint64(schema) | ||
}) | ||
schemas.get(fingerprint) match { | ||
case Some(_) => | ||
output.writeBoolean(true) | ||
output.writeLong(fingerprint) | ||
case None => | ||
output.writeBoolean(false) | ||
val compressedSchema = compress(schema) | ||
output.writeInt(compressedSchema.length) | ||
output.writeBytes(compressedSchema) | ||
} | ||
|
||
writerCache.getOrElseUpdate(schema, GenericData.get.createDatumWriter(schema)) | ||
.asInstanceOf[DatumWriter[R]] | ||
.write(datum, encoder) | ||
encoder.flush() | ||
} | ||
|
||
/** | ||
* Deserializes generic records into their in-memory form. There is internal | ||
* state to keep a cache of already seen schemas and datum readers. | ||
*/ | ||
def deserializeDatum(input: KryoInput): GenericRecord = { | ||
val schema = { | ||
if (input.readBoolean()) { | ||
val fingerprint = input.readLong() | ||
schemaCache.getOrElseUpdate(fingerprint, { | ||
schemas.get(fingerprint) match { | ||
case Some(s) => new Schema.Parser().parse(s) | ||
case None => throw new RuntimeException(s"Unknown fingerprint: $fingerprint") | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I know it'll be apparent in the stack trace, but what do you think about making the message was just a bit more clear, eg. "Error reading attempting to read avro data -- encountered an unknown fingerprint: $fingerprint, not sure what schema to use. This could happen if you registered additional schemas after starting your spark context." (I think that explanation is plausible anyway ...) There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Why There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I will change it to throw a |
||
} | ||
}) | ||
} else { | ||
val length = input.readInt() | ||
decompress(ByteBuffer.wrap(input.readBytes(length))) | ||
} | ||
} | ||
val decoder = DecoderFactory.get.directBinaryDecoder(input, null) | ||
readerCache.getOrElseUpdate(schema, GenericData.get.createDatumReader(schema)) | ||
.asInstanceOf[DatumReader[GenericRecord]] | ||
.read(null, decoder) | ||
} | ||
|
||
override def write(kryo: Kryo, output: KryoOutput, datum: GenericRecord): Unit = | ||
serializeDatum(datum, output) | ||
|
||
override def read(kryo: Kryo, input: KryoInput, datumClass: Class[GenericRecord]): GenericRecord = | ||
deserializeDatum(input) | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -23,6 +23,8 @@ import javax.annotation.Nullable | |
|
||
import scala.reflect.ClassTag | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: import order, should be below scala imports |
||
|
||
import org.apache.avro.generic.{GenericData, GenericRecord} | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: order |
||
|
||
import com.esotericsoftware.kryo.{Kryo, KryoException} | ||
import com.esotericsoftware.kryo.io.{Input => KryoInput, Output => KryoOutput} | ||
import com.esotericsoftware.kryo.serializers.{JavaSerializer => KryoJavaSerializer} | ||
|
@@ -73,6 +75,8 @@ class KryoSerializer(conf: SparkConf) | |
.split(',') | ||
.filter(!_.isEmpty) | ||
|
||
private val avroSchemas = conf.getAvroSchema | ||
|
||
def newKryoOutput(): KryoOutput = new KryoOutput(bufferSize, math.max(bufferSize, maxBufferSize)) | ||
|
||
def newKryo(): Kryo = { | ||
|
@@ -99,6 +103,9 @@ class KryoSerializer(conf: SparkConf) | |
kryo.register(classOf[HttpBroadcast[_]], new KryoJavaSerializer()) | ||
kryo.register(classOf[PythonBroadcast], new KryoJavaSerializer()) | ||
|
||
kryo.register(classOf[GenericRecord], new GenericAvroSerializer(avroSchemas)) | ||
kryo.register(classOf[GenericData.Record], new GenericAvroSerializer(avroSchemas)) | ||
|
||
try { | ||
// Use the default classloader when calling the user registrator. | ||
Thread.currentThread.setContextClassLoader(classLoader) | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,84 @@ | ||
/* | ||
* 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 org.apache.spark.serializer | ||
|
||
import java.io.{ByteArrayInputStream, ByteArrayOutputStream} | ||
import java.nio.ByteBuffer | ||
|
||
import com.esotericsoftware.kryo.io.{Output, Input} | ||
import org.apache.avro.{SchemaBuilder, Schema} | ||
import org.apache.avro.generic.GenericData.Record | ||
|
||
import org.apache.spark.{SparkFunSuite, SharedSparkContext} | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: empty line between other imports and spark imports
|
||
|
||
class GenericAvroSerializerSuite extends SparkFunSuite with SharedSparkContext { | ||
conf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer") | ||
|
||
val schema : Schema = SchemaBuilder | ||
.record("testRecord").fields() | ||
.requiredString("data") | ||
.endRecord() | ||
val record = new Record(schema) | ||
record.put("data", "test data") | ||
|
||
test("schema compression and decompression") { | ||
val genericSer = new GenericAvroSerializer(conf.getAvroSchema) | ||
assert(schema === genericSer.decompress(ByteBuffer.wrap(genericSer.compress(schema)))) | ||
} | ||
|
||
test("record serialization and deserialization") { | ||
val genericSer = new GenericAvroSerializer(conf.getAvroSchema) | ||
|
||
val outputStream = new ByteArrayOutputStream() | ||
val output = new Output(outputStream) | ||
genericSer.serializeDatum(record, output) | ||
output.flush() | ||
output.close() | ||
|
||
val input = new Input(new ByteArrayInputStream(outputStream.toByteArray)) | ||
assert(genericSer.deserializeDatum(input) === record) | ||
} | ||
|
||
test("uses schema fingerprint to decrease message size") { | ||
val genericSerFull = new GenericAvroSerializer(conf.getAvroSchema) | ||
|
||
val output = new Output(new ByteArrayOutputStream()) | ||
|
||
val beginningNormalPosition = output.total() | ||
genericSerFull.serializeDatum(record, output) | ||
output.flush() | ||
val normalLength = output.total - beginningNormalPosition | ||
|
||
conf.registerAvroSchemas(schema) | ||
val genericSerFinger = new GenericAvroSerializer(conf.getAvroSchema) | ||
val beginningFingerprintPosition = output.total() | ||
genericSerFinger.serializeDatum(record, output) | ||
val fingerprintLength = output.total - beginningFingerprintPosition | ||
|
||
assert(fingerprintLength < normalLength) | ||
} | ||
|
||
test("caches previously seen schemas") { | ||
val genericSer = new GenericAvroSerializer(conf.getAvroSchema) | ||
val compressedSchema = genericSer.compress(schema) | ||
val decompressedScheam = genericSer.decompress(ByteBuffer.wrap(compressedSchema)) | ||
|
||
assert(compressedSchema.eq(genericSer.compress(schema))) | ||
assert(decompressedScheam.eq(genericSer.decompress(ByteBuffer.wrap(compressedSchema)))) | ||
} | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
add a newline