Make schema converter configurable (#268)

* Make schema converter configurable

* Use SPI to load custom schema converter

* Add Readme

* Optimize imports

Author: kevinwallimann

README.md

@@ -153,6 +153,40 @@ val sqlSchema = new StructType(new StructField ....
 val avroSchema = SparkAvroConversions.toAvroSchema(sqlSchema, avro_schema_name, avro_schema_namespace)
 ```
 
+#### Custom data conversions
+If you would like to use custom logic to convert from Avro to Spark, you can implement the `SchemaConverter` trait.
+The custom class is loaded in ABRiS using the service provider interface (SPI), so you need to register your class in your
+`META-INF/services` resource directory. You can then configure the custom class with its short name or the fully qualified name.
+
+**Example**
+
+Custom schema converter implementation
+```scala
+package za.co.absa.abris.avro.sql
+import org.apache.avro.Schema
+import org.apache.spark.sql.types.DataType
+
+class CustomSchemaConverter extends SchemaConverter {
+  override val shortName: String = "custom"
+  override def toSqlType(avroSchema: Schema): DataType = ???
+}
+```
+
+Provider configuration file `META-INF/services/za.co.absa.abris.avro.sql.SchemaConverter`:
+```
+za.co.absa.abris.avro.sql.CustomSchemaConverter
+```
+
+Abris configuration
+```scala
+val abrisConfig = AbrisConfig
+  .fromConfluentAvro
+  .downloadReaderSchemaByLatestVersion
+  .andTopicNameStrategy("topic123")
+  .usingSchemaRegistry(registryConfig)
+  .withSchemaConverter("custom")
+```
+
 ## Multiple schemas in one topic
 The naming strategies RecordName and TopicRecordName allow for a one topic to receive different payloads, 
 i.e. payloads containing different schemas that do not have to be compatible, 

src/main/resources/META-INF/services/za.co.absa.abris.avro.sql.SchemaConverter

@@ -0,0 +1,16 @@
+#
+# Copyright 2022 ABSA Group Limited
+#
+# 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.
+#
+za.co.absa.abris.avro.sql.DefaultSchemaConverter

src/main/scala/za/co/absa/abris/avro/sql/AvroDataToCatalyst.scala

@@ -29,6 +29,7 @@ import za.co.absa.abris.avro.read.confluent.{ConfluentConstants, SchemaManagerFa
 import za.co.absa.abris.config.InternalFromAvroConfig
 
 import java.nio.ByteBuffer
+import java.util.ServiceLoader
 import scala.collection.mutable
 import scala.util.control.NonFatal
 import scala.util.{Failure, Success, Try}
@@ -39,9 +40,11 @@ private[abris] case class AvroDataToCatalyst(
   schemaRegistryConf: Option[Map[String,String]]
 ) extends UnaryExpression with ExpectsInputTypes {
 
+  private val schemaConverter = loadSchemaConverter(config.schemaConverter)
+
   override def inputTypes: Seq[BinaryType.type] = Seq(BinaryType)
 
-  override lazy val dataType: DataType = SchemaConverters.toSqlType(readerSchema).dataType
+  override lazy val dataType: DataType = schemaConverter.toSqlType(readerSchema)
 
   override def nullable: Boolean = true
 
@@ -168,4 +171,13 @@ private[abris] case class AvroDataToCatalyst(
   override protected def withNewChildInternal(newChild: Expression): Expression =
     copy(child = newChild)
 
+  private def loadSchemaConverter(nameOpt: Option[String]) = {
+    import scala.collection.JavaConverters._
+    nameOpt match {
+      case Some(name) => ServiceLoader.load(classOf[SchemaConverter]).asScala
+        .find(c => c.shortName == name || c.getClass.getName == name)
+        .getOrElse(throw new ClassNotFoundException(s"Could not find schema converter $name"))
+      case None => new DefaultSchemaConverter()
+    }
+  }
 }

src/main/scala/za/co/absa/abris/avro/sql/DefaultSchemaConverter.scala

@@ -0,0 +1,25 @@
+/*
+ * Copyright 2022 ABSA Group Limited
+ *
+ * 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 za.co.absa.abris.avro.sql
+import org.apache.avro.Schema
+import org.apache.spark.sql.avro.SchemaConverters
+import org.apache.spark.sql.types.DataType
+
+class DefaultSchemaConverter extends SchemaConverter {
+  override val shortName: String = "default"
+  override def toSqlType(avroSchema: Schema): DataType = SchemaConverters.toSqlType(avroSchema).dataType
+}

src/main/scala/za/co/absa/abris/avro/sql/SchemaConverter.scala

@@ -0,0 +1,25 @@
+/*
+ * Copyright 2022 ABSA Group Limited
+ *
+ * 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 za.co.absa.abris.avro.sql
+
+import org.apache.avro.Schema
+import org.apache.spark.sql.types.DataType
+
+trait SchemaConverter {
+  val shortName: String
+  def toSqlType(avroSchema: Schema): DataType
+}

src/main/scala/za/co/absa/abris/config/Config.scala

@@ -330,6 +330,12 @@ class FromAvroConfig private(
       schemaRegistryConf
     )
 
+  def withSchemaConverter(schemaConverter: String): FromAvroConfig =
+    new FromAvroConfig(
+      abrisConfig + (Key.SchemaConverter -> schemaConverter),
+      schemaRegistryConf
+    )
+
   def validate(): Unit = {
     if(!abrisConfig.contains(Key.ReaderSchema)) {
       throw new IllegalArgumentException(s"Missing mandatory config property ${Key.ReaderSchema}")
@@ -346,5 +352,6 @@ object FromAvroConfig {
   private[abris] object Key {
     val ReaderSchema = "readerSchema"
     val WriterSchema = "writerSchema"
+    val SchemaConverter = "schemaConverter"
   }
 }

src/main/scala/za/co/absa/abris/config/InternalFromAvroConfig.scala

@@ -27,4 +27,8 @@ private[abris] class InternalFromAvroConfig(map: Map[String, Any]) {
   val writerSchema: Option[Schema] = map
     .get(Key.WriterSchema)
     .map(s => AvroSchemaUtils.parse(s.asInstanceOf[String]))
+
+  val schemaConverter: Option[String] = map
+    .get(Key.SchemaConverter)
+    .map(_.asInstanceOf[String])
 }

src/test/resources/META-INF/services/za.co.absa.abris.avro.sql.SchemaConverter

@@ -0,0 +1,16 @@
+#
+# Copyright 2022 ABSA Group Limited
+#
+# 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.
+#
+za.co.absa.abris.avro.sql.DummySchemaConverter

src/test/scala/za/co/absa/abris/avro/sql/AvroDataToCatalystSpec.scala

@@ -16,27 +16,16 @@
 
 package za.co.absa.abris.avro.sql
 
-import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.functions.col
-import org.scalatest.flatspec.AnyFlatSpec
+import org.apache.spark.sql.types.{IntegerType, LongType, StructField, StructType}
 import org.scalatest.BeforeAndAfterEach
+import org.scalatest.flatspec.AnyFlatSpec
 import org.scalatest.matchers.should.Matchers
 import za.co.absa.abris.avro.functions._
 import za.co.absa.abris.config.{AbrisConfig, FromAvroConfig}
 import za.co.absa.abris.examples.data.generation.TestSchemas
 
 class AvroDataToCatalystSpec extends AnyFlatSpec with Matchers with BeforeAndAfterEach {
-  private val spark = SparkSession
-    .builder()
-    .appName("unitTest")
-    .master("local[2]")
-    .config("spark.driver.bindAddress", "localhost")
-    .config("spark.ui.enabled", "false")
-    .getOrCreate()
-
-  import spark.implicits._
-
-
   it should "not print schema registry configs in the spark plan" in {
     val sensitiveData = "username:password"
     val schemaString = TestSchemas.NATIVE_SIMPLE_NESTED_SCHEMA
@@ -53,4 +42,66 @@ class AvroDataToCatalystSpec extends AnyFlatSpec with Matchers with BeforeAndAft
     column.expr.toString() should not include sensitiveData
   }
 
+  it should "use the default schema converter by default" in {
+    val schemaString = TestSchemas.NATIVE_SIMPLE_NESTED_SCHEMA
+    val dummyUrl = "dummyUrl"
+    val expectedDataType = StructType(Seq(
+      StructField("int", IntegerType, nullable = false),
+      StructField("long", LongType, nullable = false)
+    ))
+
+    val fromAvroConfig = FromAvroConfig()
+      .withReaderSchema(schemaString)
+      .withSchemaRegistryConfig(Map(
+        AbrisConfig.SCHEMA_REGISTRY_URL -> dummyUrl
+      ))
+
+    val column = from_avro(col("avroBytes"), fromAvroConfig)
+    column.expr.dataType shouldBe expectedDataType
+  }
+
+  it should "use a custom schema converter identified by the short name" in {
+    val schemaString = TestSchemas.NATIVE_SIMPLE_NESTED_SCHEMA
+    val dummyUrl = "dummyUrl"
+
+    val fromAvroConfig = FromAvroConfig()
+      .withReaderSchema(schemaString)
+      .withSchemaRegistryConfig(Map(
+        AbrisConfig.SCHEMA_REGISTRY_URL -> dummyUrl
+      ))
+      .withSchemaConverter(DummySchemaConverter.name)
+
+    val column = from_avro(col("avroBytes"), fromAvroConfig)
+    column.expr.dataType shouldBe DummySchemaConverter.dataType
+  }
+
+  it should "use a custom schema converter identified by the fully qualified name" in {
+    val schemaString = TestSchemas.NATIVE_SIMPLE_NESTED_SCHEMA
+    val dummyUrl = "dummyUrl"
+
+    val fromAvroConfig = FromAvroConfig()
+      .withReaderSchema(schemaString)
+      .withSchemaRegistryConfig(Map(
+        AbrisConfig.SCHEMA_REGISTRY_URL -> dummyUrl
+      ))
+      .withSchemaConverter("za.co.absa.abris.avro.sql.DummySchemaConverter")
+
+    val column = from_avro(col("avroBytes"), fromAvroConfig)
+    column.expr.dataType shouldBe DummySchemaConverter.dataType
+  }
+
+  it should "throw an error if the specified custom schema converter does not exist" in {
+    val schemaString = TestSchemas.NATIVE_SIMPLE_NESTED_SCHEMA
+    val dummyUrl = "dummyUrl"
+
+    val fromAvroConfig = FromAvroConfig()
+      .withReaderSchema(schemaString)
+      .withSchemaRegistryConfig(Map(
+        AbrisConfig.SCHEMA_REGISTRY_URL -> dummyUrl
+      ))
+      .withSchemaConverter("nonexistent")
+
+    val ex = intercept[ClassNotFoundException](from_avro(col("avroBytes"), fromAvroConfig))
+    ex.getMessage should include ("nonexistent")
+  }
 }

src/test/scala/za/co/absa/abris/avro/sql/DummySchemaConverter.scala

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2022 ABSA Group Limited
+ *
+ * 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 za.co.absa.abris.avro.sql
+
+import org.apache.avro.Schema
+import org.apache.spark.sql.types.{DataType, LongType, StructField, StructType}
+import za.co.absa.abris.avro.sql.DummySchemaConverter._
+
+class DummySchemaConverter extends SchemaConverter {
+  override val shortName: String = name
+  override def toSqlType(avroSchema: Schema): DataType = dataType
+}
+
+object DummySchemaConverter {
+  val name: String = "dummy"
+  val dataType: DataType = StructType(Seq(StructField("long", LongType)))
+}