Java Spark是否有Map类型Encoder？求自定义Aggregator解决办法

方法1：使用Kryo序列化编码器 #

Spark的Encoders.kryo()支持序列化任意类型，包括Map，可直接在Aggregator的bufferEncoder和outputEncoder方法中返回该编码器。

示例（Scala）：

import org.apache.spark.sql.{Encoder, Encoders}
import org.apache.spark.sql.expressions.Aggregator

case class Record(category: String, count: Int)

class MapAgg extends Aggregator[Record, Map[String, Int], Map[String, Int]] {
  // 初始空Map
  override def zero: Map[String, Int] = Map.empty
  // 单条数据聚合逻辑
  override def reduce(buffer: Map[String, Int], input: Record): Map[String, Int] = {
    buffer + (input.category -> (buffer.getOrElse(input.category, 0) + input.count))
  }
  // 合并两个缓冲区
  override def merge(buffer1: Map[String, Int], buffer2: Map[String, Int]): Map[String, Int] = {
    buffer1 ++ buffer2.map { case (k, v) => k -> (buffer1.getOrElse(k, 0) + v) }
  }
  // 输出最终结果
  override def finish(reduction: Map[String, Int]): Map[String, Int] = reduction
  // 缓冲区编码器
  override def bufferEncoder: Encoder[Map[String, Int]] = Encoders.kryo[Map[String, Int]]
  // 结果输出编码器
  override def outputEncoder: Encoder[Map[String, Int]] = Encoders.kryo[Map[String, Int]]
}

方法2：将Map转换为Tuple列表后再转回 #

如果担心Kryo的性能损耗，可以先把Map拆分为(Key, Value)格式的Tuple列表，用Spark内置的Tuple编码器处理，最后通过Spark SQL函数map_from_entries转换回Map。

示例（Java）：

import org.apache.spark.sql.Encoder;
import org.apache.spark.sql.Encoders;
import org.apache.spark.sql.expressions.Aggregator;
import scala.Tuple2;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

public class MapAggJava extends Aggregator<Record, Map<String, Integer>, List<Tuple2<String, Integer>>> {
    @Override
    public Map<String, Integer> zero() {
        return new HashMap<>();
    }

    @Override
    public Map<String, Integer> reduce(Map<String, Integer> buffer, Record input) {
        buffer.put(input.getCategory(), buffer.getOrDefault(input.getCategory(), 0) + input.getCount());
        return buffer;
    }

    @Override
    public Map<String, Integer> merge(Map<String, Integer> buffer1, Map<String, Integer> buffer2) {
        buffer2.forEach((k, v) -> buffer1.put(k, buffer1.getOrDefault(k, 0) + v));
        return buffer1;
    }

    @Override
    public List<Tuple2<String, Integer>> finish(Map<String, Integer> reduction) {
        return reduction.entrySet().stream()
                .map(entry -> new Tuple2<>(entry.getKey(), entry.getValue()))
                .collect(Collectors.toList());
    }

    @Override
    public Encoder<Map<String, Integer>> bufferEncoder() {
        return Encoders.kryo(Map.class);
    }

    @Override
    public Encoder<List<Tuple2<String, Integer>>> outputEncoder() {
        return Encoders.tuple(Encoders.STRING(), Encoders.INT());
    }
}

// 使用时转换回Map：
// dataset.agg(new MapAggJava().toColumn().as("tuple_list"))
//        .selectExpr("map_from_entries(tuple_list) as result_map")

方法3：用自定义Product类包装Map #

创建一个Scala case class（自动实现Product接口）来包装Map，然后使用Encoders.product()生成对应的编码器，最后提取包装类中的Map字段即可。

示例（Scala）：

import org.apache.spark.sql.{Encoder, Encoders}
import org.apache.spark.sql.expressions.Aggregator

case class Record(category: String, count: Int)
case class MapWrapper(result: Map[String, Int])

class MapAgg extends Aggregator[Record, Map[String, Int], MapWrapper] {
  override def zero: Map[String, Int] = Map.empty
  override def reduce(buffer: Map[String, Int], input: Record): Map[String, Int] = {
    buffer + (input.category -> (buffer.getOrElse(input.category, 0) + input.count))
  }
  override def merge(buffer1: Map[String, Int], buffer2: Map[String, Int]): Map[String, Int] = {
    buffer1 ++ buffer2.map { case (k, v) => k -> (buffer1.getOrElse(k, 0) + v) }
  }
  override def finish(reduction: Map[String, Int]): MapWrapper = MapWrapper(reduction)
  override def bufferEncoder: Encoder[Map[String, Int]] = Encoders.kryo[Map[String, Int]]
  override def outputEncoder: Encoder[MapWrapper] = Encoders.product[MapWrapper]
}

// 使用时提取Map：
// dataset.agg(new MapAgg().toColumn().as("wrapped_map"))
//        .select("wrapped_map.result").as[Map[String, Int]]

额外说明 #

• Kryo编码器通用性强，但性能略低于Spark原生编码器，对性能敏感场景优先考虑后两种方案。
• 若后续升级Spark到2.3及以上版本，可直接使用Encoders.map()方法，无需上述兼容方案。

内容的提问来源于stack exchange，提问作者Nguyễn Hải Hà