Apache Beam - map only keys / values of PCollection<K<?,?>>

Only.java

package pl.pawelsz.apache.beam;

import org.apache.beam.sdk.transforms.SimpleFunction;
import org.apache.beam.sdk.values.KV;

public class Only {
  public static class Values<TK, TI, TO> extends SimpleFunction<KV<TK, TI>, KV<TK, TO>> {

    private final SimpleFunction<TI, TO> map;

    public static <Tk, Ti, To> SimpleFunction<KV<Tk, Ti>, KV<Tk, To>> with(
        SimpleFunction<Ti, To> map) {
      return new Values<Tk, Ti, To>(map);
    }

    public Values(SimpleFunction<TI, TO> map) {
      this.map = map;
    }

    @Override
    public KV<TK, TO> apply(KV<TK, TI> kv) {
      return KV.of(kv.getKey(), map.apply(kv.getValue()));
    }
  }

  public static class Keys<TKI, TKO, TV> extends SimpleFunction<KV<TKI, TV>, KV<TKO, TV>> {

    private final SimpleFunction<TKI, TKO> map;

    public static <Tki, Tko, Tv> SimpleFunction<KV<Tki, Tv>, KV<Tko, Tv>> with(
        SimpleFunction<Tki, Tko> map) {
      return new Keys<Tki, Tko, Tv>(map);
    }

    public Keys(SimpleFunction<TKI, TKO> map) {
      this.map = map;
    }

    @Override
    public KV<TKO, TV> apply(KV<TKI, TV> kv) {
      return KV.of(map.apply(kv.getKey()), kv.getValue());
    }
  }
}