如何在 PCollection<KV<String, B>> 中使用自定义编码器？

Question

我正在尝试使用自定义 Coder 以便我可以进行一些转换，但我无法让 PCollection 使用我的自定义编码器，我怀疑 (???) 是因为包裹在一个KV里面。具体来说：

Pipeline p = Pipeline.create ...
p.getCoderRegistry().registerCoder(MyClass.class, MyClassCoder.class);

...

PCollection<String> input = ...
PCollection<KV<String, MyClass>> t = input.apply(new ToKVTransform());

当我尝试运行这样的事情时，我得到一个 java.lang.ClassCastException 和一个包含 SerializableCoder 而不是我期望的 MyClassCoder 的堆栈跟踪。

[error]     at com.google.cloud.dataflow.sdk.coders.SerializableCoder.decode(SerializableCoder.java:133)
[error]     at com.google.cloud.dataflow.sdk.coders.SerializableCoder.decode(SerializableCoder.java:50)
[error]     at com.google.cloud.dataflow.sdk.coders.KvCoder.decode(KvCoder.java:95)
[error]     at com.google.cloud.dataflow.sdk.coders.KvCoder.decode(KvCoder.java:42)

我看到另一个有点相关的问题 (Using TextIO.Write with a complicated PCollection type in Google Cloud Dataflow) 的答案是将所有内容映射到字符串，并使用它在 PCollections 周围传递内容。真的推荐的方式吗？？

（注意：实际代码是在 Scala 中，但我很确定这不是 Scala <=> Java 问题，所以我在这里将其翻译成 Java。）

更新以包含 Scala 代码和更多背景：

所以这是实际的异常本身（应该在开头包括它）：

java.lang.ClassCastException: cannot assign instance of scala.collection.immutable.HashMap$SerializationProxy to field com.example.schema.Schema.keyTypes of type scala.collection.immutable.Map in instance of com.example.schema.Schema

其中 com.example.schema.Schema 是：

case class Schema(id: String, keyTypes: Map[String, Type])

最后，SchemaCoder 是：

class SchemaCoder extends com.google.cloud.dataflow.sdk.coders.CustomCoder[Schema] {
  def decode(inputStream: InputStream, context: Context): Schema = {
    val ois = new ObjectInputStream(inputStream)
    val id: String = ois.readObject().asInstanceOf[String]
    val javaMap: java.util.Map[String, Type] = ois.readObject().asInstanceOf[java.util.Map[String, Type]]
    ois.close()

    Schema(id, javaMap.asScala.toMap)
  }

  def encode(schema: Schema, outputStream: OutputStream, context: Context): Unit = {
    val baos = new ByteArrayOutputStream()
    val oos = new ObjectOutputStream(baos)
    oos.writeObject(schema.id)
    val javaMap: java.util.Map[String, Type] = schema.keyTypes.asJava
    oos.writeObject(javaMap)
    oos.close()

    val encoded = new String(Base64.encodeBase64(baos.toByteArray()))
    outputStream.write(encoded.getBytes())
  }
}

====

Edit2：ToKVTransform 实际上是这样的：

class SchemaExtractorTransform extends PTransform[PCollection[String], PCollection[Schema]] {
  class InferSchemaFromStringWithKeyFn extends DoFn[String, KV[String, Schema]] {
    override def processElement(c: DoFn[String, KV[String, Schema]]#ProcessContext): Unit = {
      val line = c.element()
      inferSchemaFromString(line)
    }
  }

  class GetFirstFn extends DoFn[KV[String, java.lang.Iterable[Schema]], Schema] {
    override def processElement(c: DoFn[KV[String, java.lang.Iterable[Schema]], Schema]#ProcessContext): Unit = {
      val idAndSchemas: KV[String, java.lang.Iterable[Schema]] = c.element()
      val it: java.util.Iterator[Schema] = idAndSchemas.getValue().iterator()
      c.output(it.next())
    }
  }

  override def apply(inputLines: PCollection[String]): PCollection[Schema] = {
    val schemasWithKey: PCollection[KV[String, Schema]] = inputLines.apply(
      ParDo.named("InferSchemas").of(new InferSchemaFromStringWithKeyFn())
    )

    val keyed: PCollection[KV[String, java.lang.Iterable[Schema]]] = schemasWithKey.apply(
      GroupByKey.create()
    )

    val schemasOnly: PCollection[Schema] = keyed.apply(
      ParDo.named("GetFirst").of(new GetFirstFn())
    )

    schemasOnly
  }
}

Answer 1

此问题不会在 Java 中重现； Scala 正在对打破 Dataflow 编码器推理的类型做一些不同的事情。要解决此问题，您可以在 PCollection 上调用 setCoder 以显式设置其编码器，例如

schemasWithKey.setCoder(KvCoder.of(StringUtf8Coder.of(), SchemaCoder.of());

这是您的代码的 Java 版本，只是为了确保它做的事情大致相同：

public static class SchemaExtractorTransform
  extends PTransform<PCollection<String>, PCollection<Schema>> {
  class InferSchemaFromStringWithKeyFn extends DoFn<String, KV<String, Schema>> {
    public void processElement(ProcessContext c) {
      c.output(KV.of(c.element(), new Schema()));
    }
  }

  class GetFirstFn extends DoFn<KV<String, java.lang.Iterable<Schema>>, Schema> {
    private static final long serialVersionUID = 0;
    public void processElement(ProcessContext c) {
      c.output(c.element().getValue().iterator().next());
    }
  }

  public PCollection<Schema> apply(PCollection<String> inputLines) {
    PCollection<KV<String, Schema>> schemasWithKey = inputLines.apply(
        ParDo.named("InferSchemas").of(new InferSchemaFromStringWithKeyFn()));

    PCollection<KV<String, java.lang.Iterable<Schema>>> keyed =
        schemasWithKey.apply(GroupByKey.<String, Schema>create());

    PCollection<Schema> schemasOnly =
        keyed.apply(ParDo.named("GetFirst").of(new GetFirstFn()));

    return schemasOnly;
  }
}

如何在 PCollection<KV<String, B>> 中使用自定义编码器？

How to use a custom Coder in a PCollection<KV<String, B>>?

google-cloud-dataflow