Modify the Dataflow Java SDK to support the combining function in PartialGroupByKey.

* Implements PGBKOp.Combiner via Combine.KeyedCombineFn.

* Modify createPartialGroupByKeyOperation to pass combiner to PartialGroupByKeyOperation when the PartialGroupByKeyInstruction contains value combining function.

* Add new constructors of PartialGroupByKeyOperation to initiate the PartialGroupByKeyOperation object with CombiningGroupingTable.

* Set SDK major version to "2" to turn on in-place combining in PGBK combiner.

----Release Notes----
[]
-------------
Created by MOE: http://code.google.com/p/moe-java
MOE_MIGRATED_REVID=89357867

Author: bo-yuan

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/DataflowPipelineRunner.java

@@ -80,7 +80,7 @@ public class DataflowPipelineRunner extends PipelineRunner<DataflowPipelineJob>
   private DataflowPipelineRunnerHooks hooks;
 
   // Environment version information
-  private static final String ENVIRONMENT_MAJOR_VERSION = "1";
+  private static final String ENVIRONMENT_MAJOR_VERSION = "2";
 
   /**
    * Construct a runner from the provided options.

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/worker/MapTaskExecutorFactory.java

@@ -16,6 +16,8 @@
 
 package com.google.cloud.dataflow.sdk.runners.worker;
 
+import static com.google.cloud.dataflow.sdk.util.Structs.getBytes;
+
 import com.google.api.services.dataflow.model.FlattenInstruction;
 import com.google.api.services.dataflow.model.InstructionInput;
 import com.google.api.services.dataflow.model.InstructionOutput;
@@ -28,9 +30,12 @@
 import com.google.cloud.dataflow.sdk.coders.Coder;
 import com.google.cloud.dataflow.sdk.coders.KvCoder;
 import com.google.cloud.dataflow.sdk.options.PipelineOptions;
+import com.google.cloud.dataflow.sdk.transforms.Combine;
 import com.google.cloud.dataflow.sdk.util.CloudObject;
 import com.google.cloud.dataflow.sdk.util.CoderUtils;
 import com.google.cloud.dataflow.sdk.util.ExecutionContext;
+import com.google.cloud.dataflow.sdk.util.PropertyNames;
+import com.google.cloud.dataflow.sdk.util.SerializableUtils;
 import com.google.cloud.dataflow.sdk.util.Serializer;
 import com.google.cloud.dataflow.sdk.util.WindowedValue;
 import com.google.cloud.dataflow.sdk.util.WindowedValue.WindowedValueCoder;
@@ -188,18 +193,64 @@ static PartialGroupByKeyOperation createPartialGroupByKeyOperation(PipelineOptio
     OutputReceiver[] receivers =
         createOutputReceivers(instruction, counterPrefix, addCounterMutator, stateSampler, 1);
 
-    PartialGroupByKeyOperation operation =
-        new PartialGroupByKeyOperation(instruction.getSystemName(),
-            new WindowingCoderGroupingKeyCreator(keyCoder),
-            new CoderSizeEstimator(WindowedValue.getValueOnlyCoder(keyCoder)),
-            new CoderSizeEstimator(valueCoder), 0.001/*sizeEstimatorSampleRate*/, PairInfo.create(),
-            receivers, counterPrefix, addCounterMutator, stateSampler);
+    PartialGroupByKeyOperation.Combiner valueCombiner = createValueCombiner(pgbk);
+
+    PartialGroupByKeyOperation operation = new PartialGroupByKeyOperation(
+        instruction.getSystemName(),
+        new WindowingCoderGroupingKeyCreator(keyCoder),
+        new CoderSizeEstimator(WindowedValue.getValueOnlyCoder(keyCoder)),
+        new CoderSizeEstimator(valueCoder), 0.001 /*sizeEstimatorSampleRate*/, valueCombiner,
+        PairInfo.create(), receivers, counterPrefix, addCounterMutator, stateSampler);
 
     attachInput(operation, pgbk.getInput(), priorOperations);
 
     return operation;
   }
 
+  static ValueCombiner createValueCombiner(PartialGroupByKeyInstruction pgbk) throws Exception {
+    if (pgbk.getValueCombiningFn() == null) {
+      return null;
+    }
+
+    Object deserializedFn = SerializableUtils.deserializeFromByteArray(
+        getBytes(CloudObject.fromSpec(pgbk.getValueCombiningFn()), PropertyNames.SERIALIZED_FN),
+        "serialized combine fn");
+    return new ValueCombiner((Combine.KeyedCombineFn) deserializedFn);
+  }
+
+  /**
+   * Implements PGBKOp.Combiner via Combine.KeyedCombineFn.
+   */
+  public static class ValueCombiner<K, VI, VA, VO>
+      implements PartialGroupByKeyOperation.Combiner<K, VI, VA, VO> {
+    private final Combine.KeyedCombineFn<K, VI, VA, VO> combineFn;
+
+    private ValueCombiner(Combine.KeyedCombineFn<K, VI, VA, VO> combineFn) {
+      this.combineFn = combineFn;
+    }
+
+    @Override
+    public VA createAccumulator(K key) {
+      return this.combineFn.createAccumulator(key);
+    }
+
+    @Override
+    public VA add(K key, VA accumulator, VI value) {
+      this.combineFn.addInput(key, accumulator, value);
+      return accumulator;
+    }
+
+    @Override
+    public VA merge(K key, Iterable<VA> accumulators) {
+      return this.combineFn.mergeAccumulators(key, accumulators);
+    }
+
+    @Override
+    public VO extract(K key, VA accumulator) {
+      return this.combineFn.extractOutput(key, accumulator);
+    }
+  }
+
   /**
    * Implements PGBKOp.PairInfo via KVs.
    */

sdk/src/main/java/com/google/cloud/dataflow/sdk/util/common/worker/PartialGroupByKeyOperation.java

@@ -133,43 +133,79 @@ public PartialGroupByKeyOperation(
       String counterPrefix,
       CounterSet.AddCounterMutator addCounterMutator,
       StateSampler stateSampler) {
+    this(operationName, groupingKeyCreator, keySizeEstimator, valueSizeEstimator, null, pairInfo,
+        receivers, counterPrefix, addCounterMutator, stateSampler);
+  }
+
+  @SuppressWarnings({"rawtypes", "unchecked"})
+  public PartialGroupByKeyOperation(
+      String operationName,
+      GroupingKeyCreator<?> groupingKeyCreator,
+      SizeEstimator<?> keySizeEstimator, SizeEstimator<?> valueSizeEstimator,
+      Combiner combineFn,
+      PairInfo pairInfo,
+      OutputReceiver[] receivers,
+      String counterPrefix,
+      CounterSet.AddCounterMutator addCounterMutator,
+      StateSampler stateSampler) {
     super(operationName, receivers, counterPrefix, addCounterMutator, stateSampler);
-    groupingTable = new BufferingGroupingTable(
-        DEFAULT_MAX_GROUPING_TABLE_BYTES, groupingKeyCreator,
-        pairInfo, keySizeEstimator, valueSizeEstimator);
+    if (combineFn == null) {
+      groupingTable = new BufferingGroupingTable(DEFAULT_MAX_GROUPING_TABLE_BYTES,
+          groupingKeyCreator, pairInfo, keySizeEstimator, valueSizeEstimator);
+    } else {
+      groupingTable = new CombiningGroupingTable(DEFAULT_MAX_GROUPING_TABLE_BYTES,
+          groupingKeyCreator, pairInfo, combineFn, keySizeEstimator, valueSizeEstimator);
+    }
   }
 
-  @SuppressWarnings("unchecked")
+  @SuppressWarnings({"rawtypes", "unchecked"})
   public PartialGroupByKeyOperation(
       String operationName,
       GroupingKeyCreator<?> groupingKeyCreator,
       SizeEstimator<?> keySizeEstimator, SizeEstimator<?> valueSizeEstimator,
       double sizeEstimatorSampleRate,
+      Combiner combineFn,
       PairInfo pairInfo,
       OutputReceiver[] receivers,
       String counterPrefix,
-      CounterSet.AddCounterMutator addCounterMutator, StateSampler stateSampler) {
+      CounterSet.AddCounterMutator addCounterMutator,
+      StateSampler stateSampler) {
     this(operationName, groupingKeyCreator,
         new SamplingSizeEstimator(keySizeEstimator, sizeEstimatorSampleRate, 1.0),
-        new SamplingSizeEstimator(valueSizeEstimator, sizeEstimatorSampleRate, 1.0),
+        new SamplingSizeEstimator(valueSizeEstimator, sizeEstimatorSampleRate, 1.0), combineFn,
         pairInfo, receivers, counterPrefix, addCounterMutator, stateSampler);
   }
 
   /** Invoked by tests. */
+  public PartialGroupByKeyOperation(GroupingKeyCreator<?> groupingKeyCreator,
+      SizeEstimator<?> keySizeEstimator, SizeEstimator<?> valueSizeEstimator, PairInfo pairInfo,
+      OutputReceiver outputReceiver, String counterPrefix,
+      CounterSet.AddCounterMutator addCounterMutator, StateSampler stateSampler) {
+    this(groupingKeyCreator,
+        keySizeEstimator, valueSizeEstimator, null, pairInfo,
+        outputReceiver,
+        counterPrefix,
+        addCounterMutator,
+        stateSampler);
+  }
+
+  /** Invoked by tests. */
+  @SuppressWarnings({"rawtypes"})
   public PartialGroupByKeyOperation(
       GroupingKeyCreator<?> groupingKeyCreator,
       SizeEstimator<?> keySizeEstimator, SizeEstimator<?> valueSizeEstimator,
+      Combiner combineFn,
       PairInfo pairInfo,
       OutputReceiver outputReceiver,
       String counterPrefix,
       CounterSet.AddCounterMutator addCounterMutator,
       StateSampler stateSampler) {
     this("PartialGroupByKeyOperation", groupingKeyCreator,
-         keySizeEstimator, valueSizeEstimator, pairInfo,
-         new OutputReceiver[]{ outputReceiver },
-         counterPrefix,
-         addCounterMutator,
-         stateSampler);
+        keySizeEstimator, valueSizeEstimator, combineFn, pairInfo,
+        new OutputReceiver[]{ outputReceiver },
+        counterPrefix,
+        addCounterMutator,
+        stateSampler);
   }
 
   @Override

sdk/src/test/java/com/google/cloud/dataflow/sdk/util/common/worker/PartialGroupByKeyOperationTest.java

@@ -128,6 +128,89 @@ public void testRunPartialGroupByKeyOperation() throws Exception {
         counterSet);
   }
 
+  @Test
+  public void testRunPartialGroupByKeyOperationWithCombiner() throws Exception {
+    Coder keyCoder = StringUtf8Coder.of();
+    Coder valueCoder = BigEndianIntegerCoder.of();
+
+    CounterSet counterSet = new CounterSet();
+    String counterPrefix = "test-";
+    StateSampler stateSampler = new StateSampler(
+        counterPrefix, counterSet.getAddCounterMutator());
+    TestReceiver receiver =
+        new TestReceiver(new ElementByteSizeObservableCoder(
+                             WindowedValue.getValueOnlyCoder(KvCoder.of(keyCoder, valueCoder))),
+            counterSet, counterPrefix);
+
+    Combiner<WindowedValue<String>, Integer, Integer, Integer> combineFn =
+        new Combiner<WindowedValue<String>, Integer, Integer, Integer>() {
+          public Integer createAccumulator(WindowedValue<String> key) {
+            return 0;
+          }
+          public Integer add(WindowedValue<String> key, Integer accumulator, Integer value) {
+            return accumulator + value;
+          }
+          public Integer merge(WindowedValue<String> key, Iterable<Integer> accumulators) {
+            Integer sum = 0;
+            for (Integer part : accumulators) {
+              sum += part;
+            }
+            return sum;
+          }
+          public Integer extract(WindowedValue<String> key, Integer accumulator) {
+            return accumulator;
+          }
+        };
+
+    PartialGroupByKeyOperation pgbkOperation =
+        new PartialGroupByKeyOperation(
+            new WindowingCoderGroupingKeyCreator(keyCoder),
+            new CoderSizeEstimator(WindowedValue.getValueOnlyCoder(keyCoder)),
+            new CoderSizeEstimator(valueCoder),
+            combineFn,
+            PairInfo.create(),
+            receiver,
+            counterPrefix,
+            counterSet.getAddCounterMutator(),
+            stateSampler);
+
+    pgbkOperation.start();
+
+    pgbkOperation.process(WindowedValue.valueInGlobalWindow(KV.of("hi", 4)));
+    pgbkOperation.process(WindowedValue.valueInGlobalWindow(KV.of("there", 5)));
+    pgbkOperation.process(WindowedValue.valueInGlobalWindow(KV.of("hi", 6)));
+    pgbkOperation.process(WindowedValue.valueInGlobalWindow(KV.of("joe", 7)));
+    pgbkOperation.process(WindowedValue.valueInGlobalWindow(KV.of("there", 8)));
+    pgbkOperation.process(WindowedValue.valueInGlobalWindow(KV.of("hi", 9)));
+
+    pgbkOperation.finish();
+
+    assertThat(receiver.outputElems,
+               IsIterableContainingInAnyOrder.<Object>containsInAnyOrder(
+                   WindowedValue.valueInGlobalWindow(KV.of("hi", 19)),
+                   WindowedValue.valueInGlobalWindow(KV.of("there", 13)),
+                   WindowedValue.valueInGlobalWindow(KV.of("joe", 7))));
+
+    // Exact counter values depend on size of encoded data.  If encoding
+    // changes, then these expected counters should change to match.
+    assertEquals(
+        new CounterSet(
+            Counter.longs("test-PartialGroupByKeyOperation-start-msecs", SUM)
+                .resetToValue(((Counter<Long>) counterSet.getExistingCounter(
+                    "test-PartialGroupByKeyOperation-start-msecs")).getAggregate(false)),
+            Counter.longs("test-PartialGroupByKeyOperation-process-msecs", SUM)
+                .resetToValue(((Counter<Long>) counterSet.getExistingCounter(
+                    "test-PartialGroupByKeyOperation-process-msecs")).getAggregate(false)),
+            Counter.longs("test-PartialGroupByKeyOperation-finish-msecs", SUM)
+                .resetToValue(((Counter<Long>) counterSet.getExistingCounter(
+                    "test-PartialGroupByKeyOperation-finish-msecs")).getAggregate(false)),
+            Counter.longs("test_receiver_out-ElementCount", SUM)
+                .resetToValue(3L),
+            Counter.longs("test_receiver_out-MeanByteCount", MEAN)
+                .resetToValue(3, 25L)),
+        counterSet);
+  }
+
   // TODO: Add tests about early flushing when the table fills.
 
   ////////////////////////////////////////////////////////////////////////////