Handle when Dataflow service tells us that values are sorted in GBK

----Release Notes----

[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=112725244

Author: lukecwik

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

@@ -21,10 +21,12 @@
 import static com.google.cloud.dataflow.sdk.runners.worker.SourceTranslationUtils.cloudProgressToReaderProgress;
 import static com.google.cloud.dataflow.sdk.runners.worker.SourceTranslationUtils.splitRequestToApproximateSplitRequest;
 import static com.google.cloud.dataflow.sdk.util.common.Counter.AggregationKind.SUM;
+import static com.google.common.base.Preconditions.checkState;
 
 import com.google.api.services.dataflow.model.ApproximateReportedProgress;
 import com.google.api.services.dataflow.model.ApproximateSplitRequest;
 import com.google.cloud.dataflow.sdk.coders.Coder;
+import com.google.cloud.dataflow.sdk.coders.Coder.Context;
 import com.google.cloud.dataflow.sdk.coders.IterableCoder;
 import com.google.cloud.dataflow.sdk.coders.KvCoder;
 import com.google.cloud.dataflow.sdk.options.PipelineOptions;
@@ -54,6 +56,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.util.Iterator;
 import java.util.concurrent.atomic.AtomicLong;
@@ -80,7 +83,8 @@ public class GroupingShuffleReader<K, V> extends NativeReader<WindowedValue<KV<K
   // Counts how many bytes were from by a given operation from a given shuffle session.
   @Nullable Counter<Long> perOperationPerDatasetBytesCounter;
   Coder<K> keyCoder;
-  Coder<V> valueCoder;
+  Coder<?> valueCoder;
+  @Nullable Coder<?> secondaryKeyCoder;
 
   public GroupingShuffleReader(
       PipelineOptions options,
@@ -90,15 +94,16 @@ public GroupingShuffleReader(
       Coder<WindowedValue<KV<K, Iterable<V>>>> coder,
       BatchModeExecutionContext executionContext,
       CounterSet.AddCounterMutator addCounterMutator,
-      String operationName)
+      String operationName,
+      boolean valuesAreSorted)
       throws Exception {
     this.shuffleReaderConfig = shuffleReaderConfig;
     this.startShufflePosition = startShufflePosition;
     this.stopShufflePosition = stopShufflePosition;
     this.executionContext = executionContext;
     this.addCounterMutator = addCounterMutator;
     this.operationName = operationName;
-    initCoder(coder);
+    initCoder(coder, valuesAreSorted);
     // We cannot initialize perOperationPerDatasetBytesCounter here, as it
     // depends on shuffleReaderConfig, which isn't populated yet.
   }
@@ -131,7 +136,8 @@ public GroupingShuffleReaderIterator<K, V> iterator() throws IOException {
         new ChunkingShuffleBatchReader(asr)));
   }
 
-  private void initCoder(Coder<WindowedValue<KV<K, Iterable<V>>>> coder) throws Exception {
+  private void initCoder(Coder<WindowedValue<KV<K, Iterable<V>>>> coder,
+      boolean valuesAreSorted) throws Exception {
     if (!(coder instanceof WindowedValueCoder)) {
       throw new Exception("unexpected kind of coder for WindowedValue: " + coder);
     }
@@ -151,7 +157,17 @@ private void initCoder(Coder<WindowedValue<KV<K, Iterable<V>>>> coder) throws Ex
           + "a key-grouping shuffle");
     }
     IterableCoder<V> iterCoder = (IterableCoder<V>) kvValueCoder;
-    this.valueCoder = iterCoder.getElemCoder();
+    if (valuesAreSorted) {
+      checkState(iterCoder.getElemCoder() instanceof KvCoder,
+          "unexpected kind of coder for elements read from a "
+          + "key-grouping value sorting shuffle: %s", iterCoder.getElemCoder());
+      @SuppressWarnings("rawtypes")
+      KvCoder<?, ?> valueKvCoder = (KvCoder) iterCoder.getElemCoder();
+      this.secondaryKeyCoder = valueKvCoder.getKeyCoder();
+      this.valueCoder = valueKvCoder.getValueCoder();
+    } else {
+      this.valueCoder = iterCoder.getElemCoder();
+    }
   }
 
   final GroupingShuffleReaderIterator<K, V> iterator(ShuffleEntryReader reader) {
@@ -390,7 +406,20 @@ public V next() {
           // notify the bytes that have been read so far.
           notifyValueReturned(currentGroupSize.getAndSet(0L));
           try {
-            return CoderUtils.decodeFromByteArray(parentReader.valueCoder, entry.getValue());
+            if (parentReader.secondaryKeyCoder != null) {
+              ByteArrayInputStream bais = new ByteArrayInputStream(entry.getSecondaryKey());
+              @SuppressWarnings("unchecked")
+              V value = (V) KV.of(
+                  // We ignore decoding the timestamp.
+                  parentReader.secondaryKeyCoder.decode(bais, Context.NESTED),
+                  CoderUtils.decodeFromByteArray(parentReader.valueCoder, entry.getValue()));
+              return value;
+            } else {
+              @SuppressWarnings("unchecked")
+              V value = (V) CoderUtils.decodeFromByteArray(parentReader.valueCoder,
+                                                           entry.getValue());
+              return value;
+            }
           } catch (IOException exn) {
             throw new RuntimeException(exn);
           }

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

@@ -17,6 +17,7 @@
 package com.google.cloud.dataflow.sdk.runners.worker;
 
 import static com.google.api.client.util.Base64.decodeBase64;
+import static com.google.cloud.dataflow.sdk.util.Structs.getBoolean;
 import static com.google.cloud.dataflow.sdk.util.Structs.getString;
 
 import com.google.cloud.dataflow.sdk.coders.Coder;
@@ -39,7 +40,6 @@
  * Creates a GroupingShuffleReader from a CloudObject spec.
  */
 public class GroupingShuffleReaderFactory implements ReaderFactory {
-
   @Override
   public NativeReader<?> create(
       CloudObject spec,
@@ -69,7 +69,8 @@ public <K, V> GroupingShuffleReader<K, V> createTyped(
           decodeBase64(getString(spec, PropertyNames.SHUFFLE_READER_CONFIG)),
           getString(spec, PropertyNames.START_SHUFFLE_POSITION, null),
           getString(spec, PropertyNames.END_SHUFFLE_POSITION, null), coder,
-          (BatchModeExecutionContext) executionContext, addCounterMutator, operationName);
+          (BatchModeExecutionContext) executionContext, addCounterMutator, operationName,
+          getBoolean(spec, PropertyNames.SORT_VALUES, false));
     }
 
     return new GroupingShuffleReader<K, V>(options,
@@ -78,7 +79,8 @@ public <K, V> GroupingShuffleReader<K, V> createTyped(
         getString(spec, PropertyNames.END_SHUFFLE_POSITION, null),
         coder,
         (BatchModeExecutionContext) executionContext,
-        addCounterMutator, operationName);
+        addCounterMutator, operationName,
+        getBoolean(spec, PropertyNames.SORT_VALUES, false));
   }
 
   /**

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

@@ -47,10 +47,11 @@ public GroupingShuffleReaderWithFaultyBytesReadCounter(
       Coder<WindowedValue<KV<K, Iterable<V>>>> coder,
       BatchModeExecutionContext executionContext,
       CounterSet.AddCounterMutator addCounterMutator,
-      String operationName)
+      String operationName,
+      boolean sortValues)
       throws Exception {
     super(options, shuffleReaderConfig, startShufflePosition, stopShufflePosition, coder,
-        executionContext, addCounterMutator, operationName);
+        executionContext, addCounterMutator, operationName, sortValues);
   }
 
   @Override

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

@@ -20,6 +20,7 @@
 
 import com.google.cloud.dataflow.sdk.coders.BigEndianLongCoder;
 import com.google.cloud.dataflow.sdk.coders.Coder;
+import com.google.cloud.dataflow.sdk.coders.Coder.Context;
 import com.google.cloud.dataflow.sdk.coders.InstantCoder;
 import com.google.cloud.dataflow.sdk.coders.KvCoder;
 import com.google.cloud.dataflow.sdk.options.DataflowWorkerHarnessOptions;
@@ -36,6 +37,7 @@
 import com.google.cloud.dataflow.sdk.values.KV;
 import com.google.common.base.Preconditions;
 
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 
 /**
@@ -220,14 +222,18 @@ public long add(WindowedValue<T> windowedElem) throws IOException {
           Object sortKey = kvValue.getKey();
           Object sortValue = kvValue.getValue();
 
-          // TODO: Need to coordinate with the
-          // GroupingShuffleReader, to make sure it knows how to
-          // reconstruct the value from the sortKeyBytes and
-          // sortValueBytes.  Right now, it doesn't know between
-          // sorting and non-sorting GBKs.
-          secondaryKeyBytes = CoderUtils.encodeToByteArray(sortKeyCoder, sortKey);
+          // Sort values by key and then timestamp so that any GroupAlsoByWindows
+          // can run more efficiently.
+          ByteArrayOutputStream baos = new ByteArrayOutputStream();
+          sortKeyCoder.encode(sortKey, baos, Context.NESTED);
+          if (!windowedElem.getTimestamp().equals(BoundedWindow.TIMESTAMP_MIN_VALUE)) {
+            // Empty timestamp suffixes sort before all other sort value keys with
+            // the same prefix. So We can omit this suffix for this common value here
+            // for efficiency and only encode when its not the minimum timestamp.
+            InstantCoder.of().encode(windowedElem.getTimestamp(), baos, Context.OUTER);
+          }
+          secondaryKeyBytes = baos.toByteArray();
           valueBytes = CoderUtils.encodeToByteArray(sortValueCoder, sortValue);
-
         } else if (groupValues) {
           // Sort values by timestamp so that GroupAlsoByWindows can run efficiently.
           if (windowedElem.getTimestamp().equals(BoundedWindow.TIMESTAMP_MIN_VALUE)) {

sdk/src/main/java/com/google/cloud/dataflow/sdk/util/PropertyNames.java

@@ -84,6 +84,7 @@ public class PropertyNames {
   public static final String SHUFFLE_KIND = "shuffle_kind";
   public static final String SHUFFLE_READER_CONFIG = "shuffle_reader_config";
   public static final String SHUFFLE_WRITER_CONFIG = "shuffle_writer_config";
+  public static final String SORT_VALUES = "sort_values";
   public static final String START_INDEX = "start_index";
   public static final String START_OFFSET = "start_offset";
   public static final String START_SHUFFLE_POSITION = "start_shuffle_position";