Use side inputs

spotify · Oct 19, 2023 · 388f1fa · 388f1fa
1 parent e543d90
commit 388f1fa
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/scio-extra/src/main/scala/com/spotify/scio/extra/voyager/syntax/SCollectionSyntax.scala b/scio-extra/src/main/scala/com/spotify/scio/extra/voyager/syntax/SCollectionSyntax.scala
@@ -100,12 +100,11 @@ class VoyagerPairSCollectionOps(
     require(!uri.exists, s"Voyager URI ${uri.value} already exists")
 
     self.transform { in =>
-      // Do not use reifyAsIterableInGlobalWindow
-      // as iterableCoder may materialize the full collection
+      val vectors = in.asIterableSideInput
       self.context
         .parallelize(Seq((): Unit))
-        .reifySideInputAsValues(in.asIterableSideInput)
-        .map { case (_, xs) =>
+        .withSideInputs(vectors)
+        .map { case (_, ctx) =>
           val indexUri = uri.value.resolve(VoyagerUri.IndexFile)
           val namesUri = uri.value.resolve(VoyagerUri.NamesFile)
           val isLocal = ScioUtil.isLocalUri(uri.value)
@@ -119,6 +118,7 @@ class VoyagerPairSCollectionOps(
             (tmpIndex, tmpNames)
           }
 
+          val xs = ctx(vectors)
           val writer =
             new VoyagerWriter(localIndex, localNames, spaceType, storageDataType, dim, ef, m)
           writer.write(xs)
@@ -130,6 +130,7 @@ class VoyagerPairSCollectionOps(
 
           uri
         }
+        .toSCollection
     }
   }