Support Tap for SMB writes (addresses #5080)

scio-core/src/main/scala/com/spotify/scio/io/Tap.scala

@@ -54,6 +54,18 @@ trait Tap[T] extends Serializable { self =>
     /** Open data set as an [[com.spotify.scio.values.SCollection SCollection]]. */
     override def open(sc: ScioContext): SCollection[U] = self.open(sc).map(f)
   }
+
+  def flatMap[U: Coder](f: T => TraversableOnce[U]): Tap[U] = new Tap[U] {
+
+    /** Parent of this Tap before [[flatMap]]. */
+    override val parent: Option[Tap[_]] = Option(self)
+
+    /** Read data set into memory. */
+    override def value: Iterator[U] = self.value.flatMap(f)
+
+    /** Open data set as an [[com.spotify.scio.values.SCollection SCollection]]. */
+    override def open(sc: ScioContext): SCollection[U] = self.open(sc).flatMap(f)
+  }
 }
 
 case object EmptyTap extends Tap[Nothing] {

scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/AvroSortedBucketIO.java

@@ -370,7 +370,7 @@ public Write<K1, K2, T> withTempDirectory(String tempDirectory) {
 
     @SuppressWarnings("unchecked")
     @Override
-    FileOperations<T> getFileOperations() {
+    public FileOperations<T> getFileOperations() {
       return getRecordClass() == null
           ? (AvroFileOperations<T>) AvroFileOperations.of(getSchema(), getCodec())
           : (AvroFileOperations<T>)

scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/JsonSortedBucketIO.java

@@ -278,7 +278,7 @@ public Write<K1, K2> withCompression(Compression compression) {
     }
 
     @Override
-    FileOperations<TableRow> getFileOperations() {
+    public FileOperations<TableRow> getFileOperations() {
       return JsonFileOperations.of(getCompression());
     }
 

scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/ParquetAvroSortedBucketIO.java

@@ -390,7 +390,7 @@ public Write<K1, K2, T> withTempDirectory(String tempDirectory) {
     }
 
     @Override
-    FileOperations<T> getFileOperations() {
+    public FileOperations<T> getFileOperations() {
       return getRecordClass() == null
           ? ParquetAvroFileOperations.of(getSchema(), getCompression(), getConfiguration())
           : ParquetAvroFileOperations.of(getRecordClass(), getCompression(), getConfiguration());

scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/SortedBucketIO.java

@@ -334,6 +334,10 @@ ResourceId getTempDirectoryOrDefault(Pipeline pipeline) {
       return FileSystems.matchNewResource(tempLocationOpt, true);
     }
 
+    public FileOperations<V> getFileOperations() {
+      return fileOperations;
+    }
+
     @Override
     public WriteResult expand(PBegin input) {
       Preconditions.checkNotNull(outputDirectory, "outputDirectory is not set");
@@ -600,7 +604,7 @@ public abstract static class Write<K1, K2, V> extends PTransform<PCollection<V>,
 
     abstract int getSorterMemoryMb();
 
-    abstract FileOperations<V> getFileOperations();
+    public abstract FileOperations<V> getFileOperations();
 
     abstract BucketMetadata<K1, K2, V> getBucketMetadata();
 

scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/TensorFlowBucketIO.java

@@ -294,7 +294,7 @@ public Write<K1, K2> withCompression(Compression compression) {
     }
 
     @Override
-    FileOperations<Example> getFileOperations() {
+    public FileOperations<Example> getFileOperations() {
       return TensorFlowFileOperations.of(getCompression());
     }
 

scio-smb/src/main/scala/com/spotify/scio/smb/SmbIO.scala

@@ -16,9 +16,16 @@
 
 package com.spotify.scio.smb
 
+import com.spotify.scio.ScioContext
 import com.spotify.scio.coders.Coder
-import com.spotify.scio.io.{KeyedIO, TapOf, TapT, TestIO}
+import com.spotify.scio.io.{ClosedTap, KeyedIO, Tap, TapOf, TapT, TestIO}
 import com.spotify.scio.util.ScioUtil
+import org.apache.beam.sdk.extensions.smb.{BucketShardId, FileOperations}
+import org.apache.beam.sdk.extensions.smb.SortedBucketSink.WriteResult
+import org.apache.beam.sdk.io.fs.ResourceId
+import org.apache.beam.sdk.values.{KV, PCollection, TupleTag}
+
+import scala.jdk.CollectionConverters._
 
 final class SmbIO[K, T](path: String, override val keyBy: T => K)(implicit
   override val keyCoder: Coder[K]
@@ -37,4 +44,21 @@ object SmbIO {
     val normalizedPaths = paths.map(p => ScioUtil.strippedPath(p) + "/").mkString(",")
     s"SortedBucketIO($normalizedPaths)"
   }
+
+  private[scio] def tap[T: Coder](
+    fileOperations: FileOperations[T],
+    writeResult: WriteResult
+  ): ScioContext => Tap[T] =
+    (sc: ScioContext) => {
+      val bucketFiles = sc
+        .wrap(
+          writeResult
+            .expand()
+            .get(new TupleTag("WrittenFiles"))
+            .asInstanceOf[PCollection[KV[BucketShardId, ResourceId]]]
+        )
+        .materialize
+
+      bucketFiles.underlying.flatMap(kv => fileOperations.iterator(kv.getValue).asScala)
+    }
 }

scio-smb/src/main/scala/com/spotify/scio/smb/SortMergeTransform.scala

@@ -18,7 +18,7 @@ package com.spotify.scio.smb
 
 import com.spotify.scio.ScioContext
 import com.spotify.scio.coders.Coder
-import com.spotify.scio.io.{ClosedTap, EmptyTap}
+import com.spotify.scio.io.{ClosedTap, EmptyTap, TapOf}
 import com.spotify.scio.testing.TestDataManager
 import com.spotify.scio.values.{SCollection, SideInput, SideInputContext}
 import org.apache.beam.sdk.extensions.smb.{SortedBucketIOUtil, SortedBucketTransform}
@@ -82,10 +82,10 @@ object SortMergeTransform {
      */
     def via(
       transformFn: (KeyType, R, SortedBucketTransform.SerializableConsumer[W]) => Unit
-    ): ClosedTap[Nothing]
+    ): ClosedTap[W]
   }
 
-  private[smb] class WriteBuilderImpl[KeyType, R, W](
+  private[smb] class WriteBuilderImpl[KeyType, R, W: Coder](
     @transient private val sc: ScioContext,
     transform: AbsCoGbkTransform[KeyType, W],
     fromResult: CoGbkResult => R
@@ -97,7 +97,7 @@ object SortMergeTransform {
 
     override def via(
       transformFn: (KeyType, R, SortedBucketTransform.SerializableConsumer[W]) => Unit
-    ): ClosedTap[Nothing] = {
+    ): ClosedTap[W] = {
       val fn = new SortedBucketTransform.TransformFn[KeyType, W]() {
         override def writeTransform(
           keyGroup: KV[KeyType, CoGbkResult],
@@ -112,8 +112,9 @@ object SortMergeTransform {
 
       val t = transform.via(fn)
       val tfName = sc.tfName(Some("sortMergeTransform"))
-      sc.applyInternal(tfName, t)
-      ClosedTap[Nothing](EmptyTap)
+      val writeResult = sc.applyInternal(tfName, t)
+
+      ClosedTap(SmbIO.tap(t.getFileOperations, writeResult).apply(sc))
     }
   }
 
@@ -130,11 +131,10 @@ object SortMergeTransform {
 
     override def via(
       transformFn: (KeyType, R, SortedBucketTransform.SerializableConsumer[W]) => Unit
-    ): ClosedTap[Nothing] = {
+    ): ClosedTap[W] = {
       val data = read.parDo(new ViaTransform(transformFn))
-      val testOutput = TestDataManager.getOutput(sc.testId.get)
-      testOutput(SortedBucketIOUtil.testId(output))(data)
-      ClosedTap[Nothing](EmptyTap)
+      TestDataManager.getOutput(sc.testId.get)(SortedBucketIOUtil.testId(output))
+      ClosedTap(TapOf[W].saveForTest(data))
     }
   }
 
@@ -159,10 +159,10 @@ object SortMergeTransform {
         SideInputContext[_],
         SortedBucketTransform.SerializableConsumer[W]
       ) => Unit
-    ): ClosedTap[Nothing]
+    ): ClosedTap[W]
   }
 
-  private[smb] class WithSideInputsWriteBuilderImpl[KeyType, R, W](
+  private[smb] class WithSideInputsWriteBuilderImpl[KeyType, R, W: Coder](
     @transient private val sc: ScioContext,
     transform: AbsCoGbkTransform[KeyType, W],
     toR: CoGbkResult => R,
@@ -175,7 +175,7 @@ object SortMergeTransform {
         SideInputContext[_],
         SortedBucketTransform.SerializableConsumer[W]
       ) => Unit
-    ): ClosedTap[Nothing] = {
+    ): ClosedTap[W] = {
       val sideViews: java.lang.Iterable[PCollectionView[_]] = sides.map(_.view).asJava
 
       val fn = new SortedBucketTransform.TransformFnWithSideInputContext[KeyType, W]() {
@@ -197,7 +197,9 @@ object SortMergeTransform {
       }
       val t = transform.via(fn, sideViews)
       sc.applyInternal(t)
-      ClosedTap[Nothing](EmptyTap)
+
+      val writeResult = sc.applyInternal(t)
+      ClosedTap(SmbIO.tap(t.getFileOperations, writeResult).apply(sc))
     }
   }
 
@@ -213,11 +215,10 @@ object SortMergeTransform {
         SideInputContext[_],
         SortedBucketTransform.SerializableConsumer[W]
       ) => Unit
-    ): ClosedTap[Nothing] = {
+    ): ClosedTap[W] = {
       val data = read.parDo(new ViaTransformWithSideOutput(transformFn))
-      val testOutput = TestDataManager.getOutput(sc.testId.get)
-      testOutput(SortedBucketIOUtil.testId(output))(data)
-      ClosedTap[Nothing](EmptyTap)
+      TestDataManager.getOutput(sc.testId.get)(SortedBucketIOUtil.testId(output))
+      ClosedTap(TapOf[W].saveForTest(data))
     }
   }
 

scio-smb/src/main/scala/com/spotify/scio/smb/syntax/SortMergeBucketSCollectionSyntax.scala

@@ -19,12 +19,16 @@ package com.spotify.scio.smb.syntax
 
 import com.spotify.scio.annotations.experimental
 import com.spotify.scio.coders.Coder
-import com.spotify.scio.io.{ClosedTap, EmptyTap}
+import com.spotify.scio.io.{ClosedTap, EmptyTap, TapOf}
+import com.spotify.scio.smb.SmbIO
 import com.spotify.scio.testing.TestDataManager
 import com.spotify.scio.values._
 import org.apache.beam.sdk.coders.KvCoder
-import org.apache.beam.sdk.extensions.smb.{SortedBucketIO, SortedBucketIOUtil}
-import org.apache.beam.sdk.values.KV
+import org.apache.beam.sdk.extensions.smb.{BucketShardId, SortedBucketIO, SortedBucketIOUtil}
+import org.apache.beam.sdk.io.fs.ResourceId
+import org.apache.beam.sdk.values.{KV, PCollection, TupleTag}
+
+import scala.jdk.CollectionConverters._
 
 trait SortMergeBucketSCollectionSyntax {
   implicit def toSortMergeBucketKeyedSCollection[K, V](
@@ -49,16 +53,20 @@ final class SortedBucketSCollection[T](private val self: SCollection[T]) {
    *   contains information about key function, bucket and shard size, etc.
    */
   @experimental
-  def saveAsSortedBucket(write: SortedBucketIO.Write[_, _, T]): ClosedTap[Nothing] = {
+  def saveAsSortedBucket(
+    write: SortedBucketIO.Write[_, _, T]
+  ): ClosedTap[T] = {
+    val beamValueCoder = self.internal.getCoder
+    implicit val valueCoder: Coder[T] = Coder.beam(beamValueCoder)
+
     if (self.context.isTest) {
-      val testOutput = TestDataManager.getOutput(self.context.testId.get)
-      testOutput[T](SortedBucketIOUtil.testId(write))(self)
+      TestDataManager.getOutput(self.context.testId.get)(SortedBucketIOUtil.testId(write))
+      ClosedTap(TapOf[T].saveForTest(self))
     } else {
-      self.applyInternal(write)
-    }
+      val writeResult = self.applyInternal(write)
 
-    // @Todo: Implement taps for metadata/bucket elements
-    ClosedTap[Nothing](EmptyTap)
+      ClosedTap(SmbIO.tap(write.getFileOperations, writeResult).apply(self.context))
+    }
   }
 }
 
@@ -82,17 +90,18 @@ final class SortedBucketPairSCollection[K, V](private val self: SCollection[KV[K
   def saveAsPreKeyedSortedBucket(
     write: SortedBucketIO.Write[K, Void, V],
     verifyKeyExtraction: Boolean = true
-  ): ClosedTap[Nothing] = {
+  ): ClosedTap[V] = {
     val beamValueCoder = self.internal.getCoder.asInstanceOf[KvCoder[K, V]].getValueCoder
+    implicit val valueCoder: Coder[V] = Coder.beam(beamValueCoder)
+
     if (self.context.isTest) {
-      implicit val valueCoder: Coder[V] = Coder.beam(beamValueCoder)
-      val testOutput = TestDataManager.getOutput(self.context.testId.get)
-      testOutput(SortedBucketIOUtil.testId(write))(self.map(_.getValue))
+      TestDataManager.getOutput(self.context.testId.get)(SortedBucketIOUtil.testId(write))
+      ClosedTap(TapOf[V].saveForTest(self.map(_.getValue)))
     } else {
-      self.applyInternal(write.onKeyedCollection(beamValueCoder, verifyKeyExtraction))
-    }
+      val writeResult =
+        self.applyInternal(write.onKeyedCollection(beamValueCoder, verifyKeyExtraction))
 
-    // @Todo: Implement taps for metadata/bucket elements
-    ClosedTap[Nothing](EmptyTap)
+      ClosedTap(SmbIO.tap(write.getFileOperations, writeResult).apply(self.context))
+    }
   }
 }