panbingkun commented on code in PR #47888:
URL: https://github.com/apache/spark/pull/47888#discussion_r1735478092
##########
connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroFunctionsSuite.scala:
##########
@@ -371,4 +372,188 @@ class AvroFunctionsSuite extends QueryTest with
SharedSparkSession {
stop = 138)))
}
}
+
+ private def serialize(record: GenericRecord, avroSchema: String):
Array[Byte] = {
+ val schema = new Schema.Parser().parse(avroSchema)
+ val datumWriter = new GenericDatumWriter[GenericRecord](schema)
+ var outputStream: ByteArrayOutputStream = null
+ var bytes: Array[Byte] = null
+ try {
+ outputStream = new ByteArrayOutputStream()
+ val encoder = EncoderFactory.get.binaryEncoder(outputStream, null)
+ datumWriter.write(record, encoder)
+ encoder.flush()
+ bytes = outputStream.toByteArray
+ } finally {
+ if (outputStream != null) {
+ outputStream.close()
+ }
+ }
+ bytes
+ }
+
+ private def deserialize(bytes: Array[Byte], avroSchema: String):
GenericRecord = {
+ val schema = new Schema.Parser().parse(avroSchema)
+ val datumReader = new GenericDatumReader[GenericRecord](schema)
+ var inputStream: SeekableByteArrayInput = null
+ var record: GenericRecord = null
+ try {
+ inputStream = new SeekableByteArrayInput(bytes)
+ val decoder = DecoderFactory.get.binaryDecoder(inputStream, null)
+ record = datumReader.read(null, decoder)
+ } finally {
+ if (inputStream != null) {
+ inputStream.close()
+ }
+ }
+ record
+ }
+
+ test("GenericRecord serialize/deserialize") {
+ val avroSchema =
+ """
+ |{
+ | "type": "record",
+ | "name": "person",
+ | "fields": [
+ | {"name": "name", "type": "string"},
+ | {"name": "age", "type": "int"},
+ | {"name": "country", "type": "string"}
+ | ]
+ |}
+ |""".stripMargin
+ val schema = new Schema.Parser().parse(avroSchema)
+ val person = new GenericRecordBuilder(schema)
+ .set("name", "spark")
+ .set("age", 18)
+ .set("country", "usa")
+ .build()
+ val bytes = serialize(person, avroSchema)
+ val readback = deserialize(bytes, avroSchema)
+ assert(person.get("name").toString === readback.get("name").toString)
+ assert(person.get("age") === readback.get("age"))
+ assert(person.get("country").toString === readback.get("country").toString)
+ }
+
+ test("use `from_avro` to read GenericRecord(stored in `array[Byte]`
datatype)") {
+ // write: Person (avro `GenericRecord`) -> binary (serialize) -> dataframe
+ // read: dataframe -> from_avro -> struct -> Person (avro `GenericRecord`)
+ val avroSchema =
+ """
+ |{
+ | "type": "record",
+ | "name": "person",
+ | "fields": [
+ | {"name": "name", "type": "string"},
+ | {"name": "age", "type": "int"},
+ | {"name": "country", "type": "string"}
+ | ]
+ |}
+ |""".stripMargin
+ val testTable = "test_avro"
+ withTable(testTable) {
+ val schema = new Schema.Parser().parse(avroSchema)
+ val person1 = new GenericRecordBuilder(schema)
+ .set("name", "sparkA")
+ .set("age", 18)
+ .set("country", "usa")
+ .build()
+ val person2 = new GenericRecordBuilder(schema)
+ .set("name", "sparkB")
+ .set("age", 19)
+ .set("country", "usb")
+ .build()
+ Seq(person1, person2)
+ .map(p => serialize(p, avroSchema))
+ .toDF("data")
+ .repartition(1)
+ .writeTo(testTable)
+ .create()
+
+ val expectedSchema = new StructType().add("data", BinaryType)
+ assert(spark.table(testTable).schema === expectedSchema)
+
+ val avroDF = sql(s"SELECT from_avro(data, '$avroSchema', map()) FROM
$testTable")
+ val readbacks = avroDF
+ .collect()
+ .map(row =>
+ new GenericRecordBuilder(schema)
+ .set("name", row.getStruct(0).getString(0))
+ .set("age", row.getStruct(0).getInt(1))
+ .set("country", row.getStruct(0).getString(2))
+ .build())
+
+ val readbackPerson1 = readbacks.head
+ assert(readbackPerson1.get(0) === person1.get(0))
+ assert(readbackPerson1.get(1).asInstanceOf[Int] ===
person1.get(1).asInstanceOf[Int])
+ assert(readbackPerson1.get(2).toString === person1.get(2))
+
+ val readbackPerson2 = readbacks(1)
+ assert(readbackPerson2.get(0).toString === person2.get(0))
+ assert(readbackPerson2.get(1).asInstanceOf[Int] ===
person2.get(1).asInstanceOf[Int])
+ assert(readbackPerson2.get(2).toString === person2.get(2))
+ }
+ }
+
+ test("use `to_avro` to read GenericRecord(stored in `struct` datatype)") {
Review Comment:
I have refactored this code and it seems clearer now. Please help review it
in your free time, thanks!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
