com.wtanaka.beam.transforms.WithKeysIdentity Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of beam Show documentation

Apache Beam Utility Code

There is a newer version: v0.0.0-128-g8bec6d4

/*
 * com.wtanaka.beam
 *
 * Copyright (C) 2017 Wesley Tanaka 
 *
 * This program is free software: you can redistribute it and/or
 * modify it under the terms of the GNU General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see
 * .
 */
package com.wtanaka.beam.transforms;

import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.WithKeys;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;

import com.wtanaka.beam.functions.Identity;

/**
 * Call WithKeys using the identity function, to convert {@code T => KV}
 */
public class WithKeysIdentity
   extends PTransform, PCollection>>
{
   private static final long serialVersionUID = 1L;

   private WithKeysIdentity()
   {
   }

   public static  PTransform, PCollection>> of()
   {
      return new WithKeysIdentity<>();
   }

   @Override
   public PCollection> expand(final PCollection input)
   {
      final Coder oldCoder = input.getCoder();
      return input
         .apply(WithKeys.of(new Identity<>()))
         .setCoder(KvCoder.of(oldCoder, oldCoder));
   }
}