[CELEBORN-1730] Support unordered KV output for Tez

client-tez/tez/src/main/java/org/apache/tez/runtime/library/common/writers/CelebornUnorderedPartitionedKVWriter.java

@@ -0,0 +1,194 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.tez.runtime.library.common.writers;
+
+import java.io.IOException;
+import java.util.Iterator;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.RawLocalFileSystem;
+import org.apache.hadoop.io.compress.CompressionCodec;
+import org.apache.hadoop.io.serializer.Serialization;
+import org.apache.hadoop.io.serializer.SerializationFactory;
+import org.apache.hadoop.io.serializer.Serializer;
+import org.apache.tez.common.counters.TaskCounter;
+import org.apache.tez.common.counters.TezCounter;
+import org.apache.tez.runtime.api.OutputContext;
+import org.apache.tez.runtime.library.api.KeyValuesWriter;
+import org.apache.tez.runtime.library.api.Partitioner;
+import org.apache.tez.runtime.library.api.TezRuntimeConfiguration;
+import org.apache.tez.runtime.library.common.ConfigUtils;
+import org.apache.tez.runtime.library.common.TezRuntimeUtils;
+import org.apache.tez.runtime.library.sort.CelebornSortBasedPusher;
+import org.apache.tez.runtime.library.utils.CodecUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.client.CelebornTezWriter;
+import org.apache.celeborn.common.CelebornConf;
+
+public class CelebornUnorderedPartitionedKVWriter extends KeyValuesWriter {
+  private static final Logger Logger =
+      LoggerFactory.getLogger(CelebornUnorderedPartitionedKVWriter.class);
+
+  protected final OutputContext outputContext;
+  protected final Configuration conf;
+  protected final RawLocalFileSystem localFs;
+  protected final Partitioner partitioner;
+  protected final Class keyClass;
+  protected final Class valClass;
+  protected final Serializer keySerializer;
+  protected final Serializer valSerializer;
+  protected final SerializationFactory serializationFactory;
+  protected final Serialization keySerialization;
+  protected final Serialization valSerialization;
+  protected final int numOutputs;
+  protected final CompressionCodec codec;
+
+  protected final TezCounter outputRecordBytesCounter;
+  protected final TezCounter outputRecordsCounter;
+  protected final TezCounter outputBytesWithOverheadCounter;
+
+  private long availableMemory;
+  private int[] numRecordsPerPartition;
+  private long[] sizePerPartition;
+  private AtomicBoolean isShutdown = new AtomicBoolean(false);
+
+  final TezRuntimeConfiguration.ReportPartitionStats reportPartitionStats;
+
+  private CelebornSortBasedPusher pusher;
+
+  public CelebornUnorderedPartitionedKVWriter(
+      OutputContext outputContext,
+      Configuration conf,
+      int numOutputs,
+      long availableMemoryBytes,
+      CelebornTezWriter celebornTezWriter,
+      CelebornConf celebornConf) {
+    this.outputContext = outputContext;
+    this.conf = conf;
+    try {
+      this.localFs = (RawLocalFileSystem) FileSystem.getLocal(conf).getRaw();
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+    this.numOutputs = numOutputs;
+
+    // k/v serialization
+    keyClass = ConfigUtils.getIntermediateOutputKeyClass(this.conf);
+    valClass = ConfigUtils.getIntermediateOutputValueClass(this.conf);
+    serializationFactory = new SerializationFactory(this.conf);
+    keySerialization = serializationFactory.getSerialization(keyClass);
+    valSerialization = serializationFactory.getSerialization(valClass);
+    keySerializer = keySerialization.getSerializer(keyClass);
+    valSerializer = valSerialization.getSerializer(valClass);
+
+    outputRecordBytesCounter = outputContext.getCounters().findCounter(TaskCounter.OUTPUT_BYTES);
+    outputRecordsCounter = outputContext.getCounters().findCounter(TaskCounter.OUTPUT_RECORDS);
+    outputBytesWithOverheadCounter =
+        outputContext.getCounters().findCounter(TaskCounter.OUTPUT_BYTES_WITH_OVERHEAD);
+
+    // stats
+    reportPartitionStats =
+        TezRuntimeConfiguration.ReportPartitionStats.fromString(
+            conf.get(
+                TezRuntimeConfiguration.TEZ_RUNTIME_REPORT_PARTITION_STATS,
+                TezRuntimeConfiguration.TEZ_RUNTIME_REPORT_PARTITION_STATS_DEFAULT));
+    sizePerPartition = (reportPartitionStats.isEnabled()) ? new long[numOutputs] : null;
+    numRecordsPerPartition = new int[numOutputs];
+
+    // compression
+    try {
+      this.codec = CodecUtils.getCodec(conf);
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+
+    Logger.info(
+        "Instantiating Partitioner: [{}]",
+        conf.get(TezRuntimeConfiguration.TEZ_RUNTIME_PARTITIONER_CLASS));
+
+    try {
+      this.partitioner = TezRuntimeUtils.instantiatePartitioner(this.conf);
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+
+    availableMemory = availableMemoryBytes;
+    // assume that there is 64MB memory to writer shuffle data
+    if (availableMemory == 0) {
+      availableMemory = 64 * 1024 * 1024;
+    }
+    pusher =
+        new CelebornSortBasedPusher(
+            keySerializer,
+            valSerializer,
+            (int) availableMemory,
+            (int) (availableMemory * 0.8),
+            null,
+            outputRecordBytesCounter,
+            outputRecordsCounter,
+            celebornTezWriter,
+            celebornConf,
+            false);
+  }
+
+  @Override
+  public void write(Object key, Iterable<Object> iterable) throws IOException {
+    Iterator<Object> it = iterable.iterator();
+    while (it.hasNext()) {
+      write(key, it.next());
+    }
+  }
+
+  @Override
+  public void write(Object key, Object value) throws IOException {
+    if (isShutdown.get()) {
+      throw new RuntimeException("Writer already closed");
+    }
+    pusher.insert(key, value, partitioner.getPartition(key, value, numOutputs));
+  }
+
+  public void close() throws IOException {
+    pusher.close();
+    isShutdown.set(true);
+    updateTezCountersAndNotify();
+  }
+
+  private void updateTezCountersAndNotify() {
+    numRecordsPerPartition = pusher.getRecordsPerPartition();
+    if (sizePerPartition != null) {
+      sizePerPartition = pusher.getBytesPerPartition();
+    }
+    outputContext.notifyProgress();
+  }
+
+  public int[] getNumRecordsPerPartition() {
+    return numRecordsPerPartition;
+  }
+
+  public boolean reportDetailedPartitionStats() {
+    return reportPartitionStats.isPrecise();
+  }
+
+  public long[] getPartitionStats() {
+    return sizePerPartition;
+  }
+}