Avro4s,如何使用自定义键类型序列化地图?

Avro4s, how to serialise a map with custom key type?

我正在使用 Avro4s。序列化 a

很容易
Map[String, T]

但我有这样的情况

sealed trait Base
case object First extends Base
case object Second extends Base

我需要连载类似

的内容
Map[Base, T]

关于实现此目标的最佳方法有什么建议吗?谢谢

事情是根据 Avro spec

Map keys are assumed to be strings.

所以 Avro 支持的唯一类型是 Map[String,T]。这意味着您需要编写一些自定义代码,将 Map[Base, T] 映射到 Map[String,T] 并返回。像这样的东西可能对你有用:

import scala.collection.breakOut
import scala.collection.immutable.Map
import scala.collection.JavaConverters._
import com.sksamuel.avro4s._
import org.apache.avro.Schema
import org.apache.avro.Schema.Field

object BaseMapAvroHelpers {
  private val nameMap: Map[Base, String] = Map(First -> "first", Second -> "second")
  private val revNameMap: Map[String, Base] = nameMap.toList.map(kv => (kv._2, kv._1)).toMap

  implicit def toSchema[T: SchemaFor]: ToSchema[Map[Base, T]] = new ToSchema[Map[Base, T]] {
    override val schema: Schema = Schema.createMap(implicitly[SchemaFor[T]].apply())
  }

  implicit def toValue[T: SchemaFor : ToValue]: ToValue[Map[Base, T]] = new ToValue[Map[Base, T]] {
    override def apply(value: Map[Base, T]): java.util.Map[String, T] = value.map(kv => (nameMap(kv._1), kv._2)).asJava
  }

  implicit def fromValue[T: SchemaFor : FromValue]: FromValue[Map[Base, T]] = new FromValue[Map[Base, T]] {
    override def apply(value: Any, field: Field): Map[Base, T] = {
      val fromValueS = implicitly[FromValue[String]]
      val fromValueT = implicitly[FromValue[T]]
      value.asInstanceOf[java.util.Map[Any, Any]].asScala.map(kv => (revNameMap(fromValueS(kv._1)), fromValueT(kv._2)))(breakOut)
    }
  }
}

用法示例:

case class Wrapper[T](value: T)
def test(): Unit = {
  import BaseMapAvroHelpers._
  val map: Map[Base, String] = Map(First -> "abc", Second -> "xyz")
  val wrapper = Wrapper(map)
  val schema = AvroSchema[Wrapper[Map[Base, String]]]
  println(s"Schema: $schema")

  val bufOut = new ByteArrayOutputStream()
  val out = AvroJsonOutputStream[Wrapper[Map[Base, String]]](bufOut)
  out.write(wrapper)
  out.flush()
  println(s"Avro Out: ${bufOut.size}")
  println(bufOut.toString("UTF-8"))

  val in = AvroJsonInputStream[Wrapper[Map[Base, String]]](new ByteArrayInputStream(bufOut.toByteArray))
  val read = in.singleEntity
  println(s"read: $read")
}

输出类似于:

Schema: {"type":"record","name":"Wrapper","namespace":"so","fields":[{"name":"value","type":{"type":"map","values":"string"}}]}
Avro Out: 40
{"value":{"first":"abc","second":"xyz"}}
read: Success(Wrapper(Map(First -> abc, Second -> xyz)))