com.google.cloud.dataflow.sdk.util.PerKeyCombineFnRunners Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of google-cloud-dataflow-java-sdk-all Show documentation
Google Cloud Dataflow Java SDK provides a simple, Java-based interface for processing virtually any size data using Google cloud resources. This artifact includes entire Dataflow Java SDK.
There is a newer version: 2.5.0
Show newest version
/*
 * Copyright (C) 2015 Google Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy of
 * the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */
package com.google.cloud.dataflow.sdk.util;

import com.google.cloud.dataflow.sdk.options.PipelineOptions;
import com.google.cloud.dataflow.sdk.transforms.Combine.KeyedCombineFn;
import com.google.cloud.dataflow.sdk.transforms.CombineFnBase.PerKeyCombineFn;
import com.google.cloud.dataflow.sdk.transforms.CombineWithContext;
import com.google.cloud.dataflow.sdk.transforms.CombineWithContext.KeyedCombineFnWithContext;
import com.google.cloud.dataflow.sdk.transforms.CombineWithContext.RequiresContextInternal;
import com.google.cloud.dataflow.sdk.transforms.DoFn;
import com.google.cloud.dataflow.sdk.transforms.windowing.BoundedWindow;
import com.google.common.collect.Iterables;

import java.util.Collection;

/**
 * Static utility methods that provide {@link PerKeyCombineFnRunner} implementations
 * for different keyed combine functions.
 */
public class PerKeyCombineFnRunners {
  /**
   * Returns a {@link PerKeyCombineFnRunner} from a {@link PerKeyCombineFn}.
   */
  public static  PerKeyCombineFnRunner
  create(PerKeyCombineFn perKeyCombineFn) {
    if (perKeyCombineFn instanceof RequiresContextInternal) {
      return new KeyedCombineFnWithContextRunner<>(
          (KeyedCombineFnWithContext) perKeyCombineFn);
    } else {
      return new KeyedCombineFnRunner<>(
          (KeyedCombineFn) perKeyCombineFn);
    }
  }

  /**
   * An implementation of {@link PerKeyCombineFnRunner} with {@link KeyedCombineFn}.
   *
   * It forwards functions calls to the {@link KeyedCombineFn}.
   */
  private static class KeyedCombineFnRunner
      implements PerKeyCombineFnRunner {
    private final KeyedCombineFn keyedCombineFn;

    private KeyedCombineFnRunner(
        KeyedCombineFn keyedCombineFn) {
      this.keyedCombineFn = keyedCombineFn;
    }

    @Override
    public KeyedCombineFn fn() {
      return keyedCombineFn;
    }

    @Override
    public AccumT createAccumulator(K key, DoFn.ProcessContext c) {
      return keyedCombineFn.createAccumulator(key);
    }

    @Override
    public AccumT addInput(
        K key, AccumT accumulator, InputT input, DoFn.ProcessContext c) {
      return keyedCombineFn.addInput(key, accumulator, input);
    }

    @Override
    public AccumT mergeAccumulators(
        K key, Iterable accumulators, DoFn.ProcessContext c) {
      return keyedCombineFn.mergeAccumulators(key, accumulators);
    }

    @Override
    public OutputT extractOutput(K key, AccumT accumulator, DoFn.ProcessContext c) {
      return keyedCombineFn.extractOutput(key, accumulator);
    }

    @Override
    public AccumT compact(K key, AccumT accumulator, DoFn.ProcessContext c) {
      return keyedCombineFn.compact(key, accumulator);
    }

    @Override
    public OutputT apply(K key, Iterable inputs, DoFn.ProcessContext c) {
      return keyedCombineFn.apply(key, inputs);
    }

    @Override
    public AccumT addInputs(K key, Iterable inputs, DoFn.ProcessContext c) {
      AccumT accum = keyedCombineFn.createAccumulator(key);
      for (InputT input : inputs) {
        accum = keyedCombineFn.addInput(key, accum, input);
      }
      return accum;
    }

    @Override
    public String toString() {
      return keyedCombineFn.toString();
    }

    @Override
    public AccumT createAccumulator(K key, PipelineOptions options,
        SideInputReader sideInputReader, Collection windows) {
      return keyedCombineFn.createAccumulator(key);
    }

    @Override
    public AccumT addInput(K key, AccumT accumulator, InputT input, PipelineOptions options,
        SideInputReader sideInputReader, Collection windows) {
      return keyedCombineFn.addInput(key, accumulator, input);
    }

    @Override
    public AccumT mergeAccumulators(K key, Iterable accumulators, PipelineOptions options,
        SideInputReader sideInputReader, Collection windows) {
      return keyedCombineFn.mergeAccumulators(key, accumulators);
    }

    @Override
    public OutputT extractOutput(K key, AccumT accumulator, PipelineOptions options,
        SideInputReader sideInputReader, Collection windows) {
      return keyedCombineFn.extractOutput(key, accumulator);
    }

    @Override
    public AccumT compact(K key, AccumT accumulator, PipelineOptions options,
        SideInputReader sideInputReader, Collection windows) {
      return keyedCombineFn.compact(key, accumulator);
    }
  }

  /**
   * An implementation of {@link PerKeyCombineFnRunner} with {@link KeyedCombineFnWithContext}.
   *
   * It forwards functions calls to the {@link KeyedCombineFnWithContext}.
   */
  private static class KeyedCombineFnWithContextRunner
      implements PerKeyCombineFnRunner {
    private final KeyedCombineFnWithContext keyedCombineFnWithContext;

    private KeyedCombineFnWithContextRunner(
        KeyedCombineFnWithContext keyedCombineFnWithContext) {
      this.keyedCombineFnWithContext = keyedCombineFnWithContext;
    }

    @Override
    public KeyedCombineFnWithContext fn() {
      return keyedCombineFnWithContext;
    }

    @Override
    public AccumT createAccumulator(K key, DoFn.ProcessContext c) {
      return keyedCombineFnWithContext.createAccumulator(key,
          CombineContextFactory.createFromProcessContext(c));
    }

    @Override
    public AccumT addInput(
        K key, AccumT accumulator, InputT value, DoFn.ProcessContext c) {
      return keyedCombineFnWithContext.addInput(key, accumulator, value,
          CombineContextFactory.createFromProcessContext(c));
    }

    @Override
    public AccumT mergeAccumulators(
        K key, Iterable accumulators, DoFn.ProcessContext c) {
      return keyedCombineFnWithContext.mergeAccumulators(
          key, accumulators, CombineContextFactory.createFromProcessContext(c));
    }

    @Override
    public OutputT extractOutput(K key, AccumT accumulator, DoFn.ProcessContext c) {
      return keyedCombineFnWithContext.extractOutput(key, accumulator,
          CombineContextFactory.createFromProcessContext(c));
    }

    @Override
    public AccumT compact(K key, AccumT accumulator, DoFn.ProcessContext c) {
      return keyedCombineFnWithContext.compact(key, accumulator,
          CombineContextFactory.createFromProcessContext(c));
    }

    @Override
    public OutputT apply(K key, Iterable inputs, DoFn.ProcessContext c) {
      return keyedCombineFnWithContext.apply(key, inputs,
          CombineContextFactory.createFromProcessContext(c));
    }

    @Override
    public AccumT addInputs(K key, Iterable inputs, DoFn.ProcessContext c) {
      CombineWithContext.Context combineContext = CombineContextFactory.createFromProcessContext(c);
      AccumT accum = keyedCombineFnWithContext.createAccumulator(key, combineContext);
      for (InputT input : inputs) {
        accum = keyedCombineFnWithContext.addInput(key, accum, input, combineContext);
      }
      return accum;
    }

    @Override
    public String toString() {
      return keyedCombineFnWithContext.toString();
    }

    @Override
    public AccumT createAccumulator(K key, PipelineOptions options, SideInputReader sideInputReader,
        Collection windows) {
      return keyedCombineFnWithContext.createAccumulator(key,
        CombineContextFactory.createFromComponents(
          options, sideInputReader, Iterables.getOnlyElement(windows)));
    }

    @Override
    public AccumT addInput(K key, AccumT accumulator, InputT input, PipelineOptions options,
        SideInputReader sideInputReader, Collection windows) {
      return keyedCombineFnWithContext.addInput(key, accumulator, input,
        CombineContextFactory.createFromComponents(
          options, sideInputReader, Iterables.getOnlyElement(windows)));
    }

    @Override
    public AccumT mergeAccumulators(K key, Iterable accumulators, PipelineOptions options,
        SideInputReader sideInputReader, Collection windows) {
      return keyedCombineFnWithContext.mergeAccumulators(key, accumulators,
        CombineContextFactory.createFromComponents(
          options, sideInputReader, Iterables.getOnlyElement(windows)));
    }

    @Override
    public OutputT extractOutput(K key, AccumT accumulator, PipelineOptions options,
        SideInputReader sideInputReader, Collection windows) {
      return keyedCombineFnWithContext.extractOutput(key, accumulator,
        CombineContextFactory.createFromComponents(
          options, sideInputReader, Iterables.getOnlyElement(windows)));
    }

    @Override
    public AccumT compact(K key, AccumT accumulator, PipelineOptions options,
        SideInputReader sideInputReader, Collection windows) {
      return keyedCombineFnWithContext.compact(key, accumulator,
        CombineContextFactory.createFromComponents(
          options, sideInputReader, Iterables.getOnlyElement(windows)));
    }
  }
}