@@ -22,24 +22,25 @@
 import com.google.cloud.dataflow.sdk.coders.Coder;
 import com.google.cloud.dataflow.sdk.coders.VoidCoder;
 import com.google.cloud.dataflow.sdk.io.Read.Bounded;
-import com.google.cloud.dataflow.sdk.runners.DirectPipelineRunner;
+import com.google.cloud.dataflow.sdk.options.PipelineOptions;
 import com.google.cloud.dataflow.sdk.runners.PipelineRunner;
-import com.google.cloud.dataflow.sdk.runners.worker.AvroSink;
 import com.google.cloud.dataflow.sdk.transforms.PTransform;
 import com.google.cloud.dataflow.sdk.util.IOChannelUtils;
-import com.google.cloud.dataflow.sdk.util.WindowedValue;
-import com.google.cloud.dataflow.sdk.util.common.worker.Sink;
+import com.google.cloud.dataflow.sdk.util.MimeTypes;
 import com.google.cloud.dataflow.sdk.values.PCollection;
 import com.google.cloud.dataflow.sdk.values.PDone;
 import com.google.cloud.dataflow.sdk.values.PInput;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 
 import org.apache.avro.Schema;
+import org.apache.avro.file.DataFileWriter;
 import org.apache.avro.generic.GenericRecord;
 import org.apache.avro.reflect.ReflectData;
 
 import java.io.IOException;
-import java.util.List;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
 import java.util.regex.Pattern;
 
 import javax.annotation.Nullable;
@@ -317,7 +318,7 @@ public PCollection<T> apply(PInput input) {
                 : com.google.cloud.dataflow.sdk.io.Read.from(
                     AvroSource.from(filepattern).withSchema(type));
 
-        PCollection<T> pcol = input.getPipeline().apply(read);
+        PCollection<T> pcol = input.getPipeline().apply("Read", read);
         // Honor the default output coder that would have been used by this PTransform.
         pcol.setCoder(getDefaultOutputCoder());
         return pcol;
@@ -473,8 +474,6 @@ public static class Bound<T> extends PTransform<PCollection<T>, PDone> {
       final int numShards;
       /** Shard template string. */
       final String shardTemplate;
-      /** Insert a shuffle before writing to decouple parallelism when numShards != 0. */
-      final boolean forceReshard;
       /** The class type of the records. */
       final Class<T> type;
       /** The schema of the output file. */
@@ -484,18 +483,23 @@ public static class Bound<T> extends PTransform<PCollection<T>, PDone> {
       final boolean validate;
 
       Bound(Class<T> type) {
-        this(null, null, "", 0, ShardNameTemplate.INDEX_OF_MAX, true, type, null, true);
+        this(null, null, "", 0, ShardNameTemplate.INDEX_OF_MAX, type, null, true);
       }
 
-      Bound(String name, String filenamePrefix, String filenameSuffix, int numShards,
-          String shardTemplate, boolean forceReshard, Class<T> type, Schema schema,
+      Bound(
+          String name,
+          String filenamePrefix,
+          String filenameSuffix,
+          int numShards,
+          String shardTemplate,
+          Class<T> type,
+          Schema schema,
           boolean validate) {
         super(name);
         this.filenamePrefix = filenamePrefix;
         this.filenameSuffix = filenameSuffix;
         this.numShards = numShards;
         this.shardTemplate = shardTemplate;
-        this.forceReshard = forceReshard;
         this.type = type;
         this.schema = schema;
         this.validate = validate;
@@ -509,8 +513,7 @@ public static class Bound<T> extends PTransform<PCollection<T>, PDone> {
        */
       public Bound<T> named(String name) {
         return new Bound<>(
-            name, filenamePrefix, filenameSuffix, numShards, shardTemplate, forceReshard,
-            type, schema, validate);
+            name, filenamePrefix, filenameSuffix, numShards, shardTemplate, type, schema, validate);
       }
 
       /**
@@ -525,8 +528,7 @@ public Bound<T> named(String name) {
       public Bound<T> to(String filenamePrefix) {
         validateOutputComponent(filenamePrefix);
         return new Bound<>(
-            name, filenamePrefix, filenameSuffix, numShards, shardTemplate, forceReshard,
-            type, schema, validate);
+            name, filenamePrefix, filenameSuffix, numShards, shardTemplate, type, schema, validate);
       }
 
       /**
@@ -540,8 +542,7 @@ public Bound<T> to(String filenamePrefix) {
       public Bound<T> withSuffix(String filenameSuffix) {
         validateOutputComponent(filenameSuffix);
         return new Bound<>(
-            name, filenamePrefix, filenameSuffix, numShards, shardTemplate, forceReshard, type,
-            schema, validate);
+            name, filenamePrefix, filenameSuffix, numShards, shardTemplate, type, schema, validate);
       }
 
       /**
@@ -559,32 +560,9 @@ public Bound<T> withSuffix(String filenameSuffix) {
        * @see ShardNameTemplate
        */
       public Bound<T> withNumShards(int numShards) {
-        return withNumShards(numShards, forceReshard);
-      }
-
-      /**
-       * Returns a new {@link PTransform} that's like this one but
-       * that uses the provided shard count.
-       *
-       * <p>Constraining the number of shards is likely to reduce
-       * the performance of a pipeline. If forceReshard is true, the output
-       * will be shuffled to obtain the desired sharding. If it is false,
-       * data will not be reshuffled, but parallelism of preceeding stages
-       * may be constrained. Setting this value is not recommended
-       * unless you require a specific number of output files.
-       *
-       * <p>Does not modify this object.
-       *
-       * @param numShards the number of shards to use, or 0 to let the system
-       *                  decide.
-       * @param forceReshard whether to force a reshard to obtain the desired sharding.
-       * @see ShardNameTemplate
-       */
-      private Bound<T> withNumShards(int numShards, boolean forceReshard) {
         Preconditions.checkArgument(numShards >= 0);
         return new Bound<>(
-            name, filenamePrefix, filenameSuffix, numShards, shardTemplate, forceReshard,
-            type, schema, validate);
+            name, filenamePrefix, filenameSuffix, numShards, shardTemplate, type, schema, validate);
       }
 
       /**
@@ -597,8 +575,7 @@ private Bound<T> withNumShards(int numShards, boolean forceReshard) {
        */
       public Bound<T> withShardNameTemplate(String shardTemplate) {
         return new Bound<>(
-            name, filenamePrefix, filenameSuffix, numShards, shardTemplate, forceReshard,
-            type, schema, validate);
+            name, filenamePrefix, filenameSuffix, numShards, shardTemplate, type, schema, validate);
       }
 
       /**
@@ -611,23 +588,7 @@ public Bound<T> withShardNameTemplate(String shardTemplate) {
        * <p>Does not modify this object.
        */
       public Bound<T> withoutSharding() {
-        return withoutSharding(forceReshard);
-      }
-
-      /**
-       * Returns a new {@link PTransform} that's like this one but
-       * that forces a single file as output.
-       *
-       * <p>This is a shortcut for
-       * {@code .withNumShards(1, forceReshard).withShardNameTemplate("")}
-       *
-       * <p>Does not modify this object.
-       *
-       * @param forceReshard whether to force a reshard to obtain the desired sharding.
-       */
-      private Bound<T> withoutSharding(boolean forceReshard) {
-        return new Bound<>(name, filenamePrefix, filenameSuffix, 1, "", forceReshard,
-            type, schema, validate);
+        return new Bound<>(name, filenamePrefix, filenameSuffix, 1, "", type, schema, validate);
       }
 
       /**
@@ -640,8 +601,15 @@ private Bound<T> withoutSharding(boolean forceReshard) {
        * @param <X> the type of the elements of the input PCollection
        */
       public <X> Bound<X> withSchema(Class<X> type) {
-        return new Bound<>(name, filenamePrefix, filenameSuffix, numShards, shardTemplate,
-            forceReshard, type, ReflectData.get().getSchema(type), validate);
+        return new Bound<>(
+            name,
+            filenamePrefix,
+            filenameSuffix,
+            numShards,
+            shardTemplate,
+            type,
+            ReflectData.get().getSchema(type),
+            validate);
       }
 
       /**
@@ -652,8 +620,15 @@ public <X> Bound<X> withSchema(Class<X> type) {
        * <p>Does not modify this object.
        */
       public Bound<GenericRecord> withSchema(Schema schema) {
-        return new Bound<>(name, filenamePrefix, filenameSuffix, numShards, shardTemplate,
-            forceReshard, GenericRecord.class, schema, validate);
+        return new Bound<>(
+            name,
+            filenamePrefix,
+            filenameSuffix,
+            numShards,
+            shardTemplate,
+            GenericRecord.class,
+            schema,
+            validate);
       }
 
       /**
@@ -679,8 +654,7 @@ public Bound<GenericRecord> withSchema(String schema) {
        */
       public Bound<T> withoutValidation() {
         return new Bound<>(
-            name, filenamePrefix, filenameSuffix, numShards, shardTemplate, forceReshard,
-            type, schema, false);
+            name, filenamePrefix, filenameSuffix, numShards, shardTemplate, type, schema, false);
       }
 
       @Override
@@ -693,14 +667,14 @@ public PDone apply(PCollection<T> input) {
           throw new IllegalStateException("need to set the schema of an AvroIO.Write transform");
         }
 
-        if (numShards > 0 && forceReshard) {
-          // Reshard and re-apply a version of this write without resharding.
-          return input
-              .apply(new FileBasedSink.ReshardForWrite<T>())
-              .apply(withNumShards(numShards, false));
-        } else {
-          return PDone.in(input.getPipeline());
-        }
+        // Note that custom sinks currently do not expose sharding controls.
+        // Thus pipeline runner writers need to individually add support internally to
+        // apply user requested sharding limits.
+        return input.apply(
+            "Write",
+            com.google.cloud.dataflow.sdk.io.Write.to(
+                new AvroSink<>(
+                    filenamePrefix, filenameSuffix, shardTemplate, AvroCoder.of(type, schema))));
       }
 
       /**
@@ -742,21 +716,6 @@ public Schema getSchema() {
       public boolean needsValidation() {
         return validate;
       }
-
-      static {
-        @SuppressWarnings("rawtypes")
-        DirectPipelineRunner.TransformEvaluator<Bound> transformEvaluator =
-            new DirectPipelineRunner.TransformEvaluator<Bound>() {
-          @Override
-          @SuppressWarnings("unchecked")
-          public void evaluate(
-              Bound transform, DirectPipelineRunner.EvaluationContext context) {
-            evaluateWriteHelper(transform, context);
-          }
-        };
-        DirectPipelineRunner.registerDefaultTransformEvaluator(
-            Bound.class, transformEvaluator);
-      }
     }
 
     /** Disallow construction of utility class. */
@@ -779,25 +738,73 @@ private static void validateOutputComponent(String partialFilePattern) {
   /** Disallow construction of utility class. */
   private AvroIO() {}
 
-  private static <T> void evaluateWriteHelper(
-      Write.Bound<T> transform, DirectPipelineRunner.EvaluationContext context) {
-    List<WindowedValue<T>> elems =
-        context.getPCollectionWindowedValues(context.getInput(transform));
-    int numShards = transform.numShards;
-    if (numShards < 1) {
-      // System gets to choose. For direct mode, choose 1.
-      numShards = 1;
+  /**
+   * A {@link FileBasedSink} for Avro files.
+   */
+  @VisibleForTesting
+  static class AvroSink<T> extends FileBasedSink<T> {
+    private final AvroCoder<T> coder;
+
+    @VisibleForTesting
+    AvroSink(
+        String baseOutputFilename, String extension, String fileNameTemplate, AvroCoder<T> coder) {
+      super(baseOutputFilename, extension, fileNameTemplate);
+      this.coder = coder;
+    }
+
+    @Override
+    public FileBasedSink.FileBasedWriteOperation<T> createWriteOperation(PipelineOptions options) {
+      return new AvroWriteOperation<>(this, coder);
+    }
+
+    /**
+     * A {@link com.google.cloud.dataflow.sdk.io.FileBasedSink.FileBasedWriteOperation
+     * FileBasedWriteOperation} for Avro files.
+     */
+    private static class AvroWriteOperation<T> extends FileBasedWriteOperation<T> {
+      private final AvroCoder<T> coder;
+
+      private AvroWriteOperation(AvroSink<T> sink, AvroCoder<T> coder) {
+        super(sink);
+        this.coder = coder;
+      }
+
+      @Override
+      public FileBasedWriter<T> createWriter(PipelineOptions options) throws Exception {
+        return new AvroWriter<>(this, coder);
+      }
     }
-    AvroSink<T> writer = new AvroSink<>(transform.filenamePrefix, transform.shardTemplate,
-        transform.filenameSuffix, numShards,
-        WindowedValue.getValueOnlyCoder(AvroCoder.of(transform.type, transform.schema)));
-    try (Sink.SinkWriter<WindowedValue<T>> sink = writer.writer()) {
-      for (WindowedValue<T> elem : elems) {
-        sink.add(elem);
-      }
-    } catch (IOException exn) {
-      throw new RuntimeException(
-          "unable to write to output file \"" + transform.filenamePrefix + "\"", exn);
+
+    /**
+     * A {@link com.google.cloud.dataflow.sdk.io.FileBasedSink.FileBasedWriter FileBasedWriter}
+     * for Avro files.
+     */
+    private static class AvroWriter<T> extends FileBasedWriter<T> {
+      private final AvroCoder<T> coder;
+      private DataFileWriter<T> dataFileWriter;
+
+      public AvroWriter(FileBasedWriteOperation<T> writeOperation, AvroCoder<T> coder) {
+        super(writeOperation);
+        this.mimeType = MimeTypes.BINARY;
+        this.coder = coder;
+      }
+
+      @SuppressWarnings("deprecation") // uses internal test functionality.
+      @Override
+      protected void prepareWrite(WritableByteChannel channel) throws Exception {
+        dataFileWriter = new DataFileWriter<>(coder.createDatumWriter());
+        dataFileWriter.create(coder.getSchema(), Channels.newOutputStream(channel));
+      }
+
+      @Override
+      public void write(T value) throws Exception {
+        dataFileWriter.append(value);
+      }
+
+      @Override
+      protected void writeFooter() throws Exception {
+        dataFileWriter.flush();
+      }
     }
   }
 }