forklift SMB implementation into scio-smb module

build.sbt

@@ -31,6 +31,7 @@ val annoy4sVersion = "0.9.0"
 val annoyVersion = "0.2.6"
 val asmVersion = "4.13"
 val autoServiceVersion = "1.0-rc2"
+val autoValueVersion = "1.7"
 val avroVersion = "1.8.2"
 val breezeVersion = "1.0"
 val chillVersion = "0.9.4"
@@ -364,7 +365,8 @@ lazy val root: Project = Project("scio", file("."))
     `scio-examples`,
     `scio-repl`,
     `scio-jmh`,
-    `scio-macros`
+    `scio-macros`,
+    `scio-smb`
   )
 
 lazy val `scio-core`: Project = project
@@ -979,6 +981,34 @@ lazy val `scio-jmh`: Project = project
   )
   .enablePlugins(JmhPlugin)
 
+lazy val `scio-smb`: Project = project
+  .in(file("scio-smb"))
+  .settings(commonSettings)
+  .settings(itSettings)
+  .settings(beamRunnerSettings)
+  .settings(
+    description := "Sort Merge Bucket source/sink implementations for Apache Beam",
+    libraryDependencies ++= Seq(
+      "org.apache.beam" % "beam-sdks-java-core" % beamVersion,
+      "org.apache.beam" % "beam-sdks-java-core" % "it,test" classifier "tests",
+      "org.apache.beam" % "beam-sdks-java-extensions-sorter" % beamVersion,
+      "org.apache.beam" % "beam-sdks-java-extensions-protobuf" % beamVersion,
+      "org.apache.beam" % "beam-sdks-java-io-google-cloud-platform" % beamVersion,
+      "com.google.apis" % "google-api-services-bigquery" % googleApiServicesBigQuery,
+      "org.tensorflow" % "proto" % tensorFlowVersion,
+      "com.google.auto.value" % "auto-value-annotations" % autoValueVersion,
+      "com.google.auto.value" % "auto-value" % autoValueVersion,
+      "javax.annotation" % "javax.annotation-api" % "1.3.2",
+      "org.hamcrest" % "hamcrest-all" % hamcrestVersion % Test,
+      "com.novocode" % "junit-interface" % "0.11" % Test,
+      "junit" % "junit" % "4.13-beta-1" % Test
+    ),
+    Test / classLoaderLayeringStrategy := ClassLoaderLayeringStrategy.Flat
+  )
+  .configs(
+    IntegrationTest
+  )
+
 lazy val site: Project = project
   .in(file("site"))
   .settings(commonSettings)

scio-core/src/main/scala/com/spotify/scio/ScioContext.scala

@@ -428,6 +428,10 @@ class ScioContext private[scio] (
   val options: PipelineOptions,
   private var artifacts: List[String]
 ) extends TransformNameable {
+  // var _pipeline member is lazily initialized, this makes sure that file systems are registered
+  // before any IO
+  FileSystems.setDefaultPipelineOptions(options)
+
   /** Get PipelineOptions as a more specific sub-type. */
   def optionsAs[T <: PipelineOptions: ClassTag]: T =
     options.as(ScioUtil.classOf[T])

scio-smb/src/it/java/org/apache/beam/sdk/extensions/smb/SmbIT.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright 2019 Spotify AB.
+ *
+ * Licensed 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.beam.sdk.extensions.smb;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Comparator;
+import org.apache.beam.sdk.extensions.smb.benchmark.CoGroupByKeyBenchmark;
+import org.apache.beam.sdk.extensions.smb.benchmark.SinkBenchmark;
+import org.apache.beam.sdk.extensions.smb.benchmark.SourceBenchmark;
+
+/** Integration test intended for local execution with DirectRunner. */
+public class SmbIT {
+
+  public static void main(String[] args) throws IOException {
+    final Path temp = Files.createTempDirectory("smb-");
+
+    final Path avroSource = temp.resolve("avro");
+    final Path jsonSource = temp.resolve("json");
+    final Path tempLocation = temp.resolve("temp");
+
+    final String[] smbSinkArgs = new String[args.length + 3];
+    smbSinkArgs[0] = "--avroDestination=" + avroSource;
+    smbSinkArgs[1] = "--jsonDestination=" + jsonSource;
+    smbSinkArgs[2] = "--tempLocation=" + tempLocation;
+
+    System.arraycopy(args, 0, smbSinkArgs, 3, args.length);
+
+    final String[] smbSourceArgs = new String[] {
+        "--avroSource=" + avroSource,
+        "--jsonSource=" + jsonSource,
+        "--tempLocation=" + tempLocation
+    };
+
+    final String[] coGroupByKeyArgs = new String[] {
+        "--avroSource=" + avroSource.resolve("bucket-*.avro"),
+        "--jsonSource=" + jsonSource.resolve("bucket-*.json"),
+        "--tempLocation=" + tempLocation
+    };
+
+    try {
+      SinkBenchmark.main(smbSinkArgs);
+      SourceBenchmark.main(smbSourceArgs);
+
+      // Baseline comparison with default CGBK implementation
+      CoGroupByKeyBenchmark.main(coGroupByKeyArgs);
+    } finally {
+      Files.walk(temp)
+        .sorted(Comparator.reverseOrder())
+        .map(Path::toFile)
+        .forEach(File::delete);
+    }
+  }
+}

scio-smb/src/it/java/org/apache/beam/sdk/extensions/smb/SmbPublicAPITest.java

@@ -0,0 +1,104 @@
+/*
+ * Copyright 2019 Spotify AB.
+ *
+ * Licensed 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.beam.sdk.extensions.smb;
+
+import com.google.api.services.bigquery.model.TableRow;
+import java.io.IOException;
+import java.util.Collections;
+import org.apache.beam.sdk.coders.CannotProvideCoderException;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.extensions.smb.BucketMetadata.HashType;
+import org.apache.beam.sdk.extensions.smb.SortedBucketSource.BucketedInput;
+import org.apache.beam.sdk.io.AvroGeneratedUser;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.util.MimeTypes;
+import org.apache.beam.sdk.values.TupleTag;
+import org.tensorflow.example.Example;
+
+/** Test public API access level. Passes by successfully compiling. */
+public class SmbPublicAPITest {
+  public static void main(String[] args)
+      throws Coder.NonDeterministicException, CannotProvideCoderException, IOException {
+    // public API
+    AvroSortedBucketIO.write(String.class, "name", AvroGeneratedUser.class).to("avro");
+    JsonSortedBucketIO.write(String.class, "name").to("json");
+    TensorFlowBucketIO.write(String.class, "name").to("tf");
+
+    TupleTag<AvroGeneratedUser> avro = new TupleTag<>("avro");
+    TupleTag<TableRow> json = new TupleTag<>("json");
+    TupleTag<Example> tf = new TupleTag<>("tf");
+    SortedBucketIO.read(String.class)
+        .of(AvroSortedBucketIO.read(avro, AvroGeneratedUser.class).from("avro"))
+        .and(JsonSortedBucketIO.read(json).from("json"))
+        .and(TensorFlowBucketIO.read(tf).from("tf"));
+
+    // extendable API
+    new SortedBucketSink<>(
+        new MyMetadata(8, 1, String.class, HashType.MURMUR3_32),
+        FileSystems.matchNewResource("output", true),
+        FileSystems.matchNewResource("temp", true),
+        ".avro",
+        new MyFileOperation(),
+        1);
+
+    new SortedBucketSource<>(
+        String.class,
+        Collections.singletonList(
+            new BucketedInput<>(
+                new TupleTag<>(),
+                FileSystems.matchSingleFileSpec("in").resourceId(),
+                ".avro",
+                new MyFileOperation())));
+  }
+
+  private static class MyMetadata extends BucketMetadata<String, String> {
+    private MyMetadata(int numBuckets, int numShards, Class<String> keyClass, BucketMetadata.HashType hashType)
+        throws CannotProvideCoderException, Coder.NonDeterministicException {
+      super(BucketMetadata.CURRENT_VERSION, numBuckets, numShards, keyClass, hashType);
+    }
+
+    @Override
+    public String extractKey(String value) {
+      return null;
+    }
+  }
+
+  private static class MyFileOperation extends FileOperations<String> {
+
+    private MyFileOperation() {
+      super(Compression.UNCOMPRESSED, MimeTypes.BINARY);
+    }
+
+    @Override
+    protected Reader<String> createReader() {
+      return null;
+    }
+
+    @Override
+    protected FileIO.Sink<String> createSink() {
+      return null;
+    }
+
+    @Override
+    public Coder<String> getCoder() {
+      return null;
+    }
+  }
+}

scio-smb/src/it/java/org/apache/beam/sdk/extensions/smb/benchmark/CoGroupByKeyBenchmark.java

@@ -0,0 +1,136 @@
+/*
+ * Copyright 2019 Spotify AB.
+ *
+ * Licensed 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.beam.sdk.extensions.smb.benchmark;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.api.services.bigquery.model.TableRow;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.PipelineResult.State;
+import org.apache.beam.sdk.coders.AvroCoder;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.io.AvroGeneratedUser;
+import org.apache.beam.sdk.io.AvroIO;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.transforms.Count;
+import org.apache.beam.sdk.transforms.FlatMapElements;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.WithKeys;
+import org.apache.beam.sdk.transforms.join.CoGroupByKey;
+import org.apache.beam.sdk.transforms.join.KeyedPCollectionTuple;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.sdk.values.TypeDescriptors;
+
+/** Benchmark of {@link CoGroupByKey} using data generated by {@link SinkBenchmark}. */
+public class CoGroupByKeyBenchmark {
+
+  private static ObjectMapper objectMapper = new ObjectMapper();
+
+  /** SourceOptions. */
+  public interface SourceOptions extends PipelineOptions {
+    String getAvroSource();
+
+    void setAvroSource(String value);
+
+    String getJsonSource();
+
+    void setJsonSource(String value);
+  }
+
+  public static void main(String[] args) {
+    final SourceOptions sourceOptions =
+        PipelineOptionsFactory.fromArgs(args).as(SourceOptions.class);
+    System.out.println("CoGroupByKey SourceOptions=" + sourceOptions);
+
+    final Pipeline pipeline = Pipeline.create(sourceOptions);
+
+    final PCollection<KV<String, AvroGeneratedUser>> lhs =
+        pipeline
+            .apply(AvroIO.read(AvroGeneratedUser.class).from(sourceOptions.getAvroSource()))
+            .apply(WithKeys.of(user -> user.getName().toString()))
+            .setCoder(KvCoder.of(StringUtf8Coder.of(), AvroCoder.of(AvroGeneratedUser.class)));
+
+    final PCollection<KV<String, TableRow>> rhs =
+        pipeline
+            .apply(TextIO.read().from(sourceOptions.getJsonSource()))
+            .apply(
+                MapElements.into(
+                        TypeDescriptors.kvs(
+                            TypeDescriptors.strings(), TypeDescriptor.of(TableRow.class)))
+                    .via(
+                        s -> {
+                          try {
+                            TableRow record = objectMapper.readValue(s, TableRow.class);
+                            return KV.of(record.get("user").toString(), record);
+                          } catch (IOException e) {
+                            throw new RuntimeException(e);
+                          }
+                        }));
+
+    final TupleTag<AvroGeneratedUser> tl = new TupleTag<>();
+    final TupleTag<TableRow> tr = new TupleTag<>();
+
+
+    KeyedPCollectionTuple.of(tl, lhs)
+        .and(tr, rhs)
+        .apply(CoGroupByKey.create())
+        .apply(
+            FlatMapElements.into(
+                    TypeDescriptors.kvs(
+                        TypeDescriptors.strings(),
+                        TypeDescriptors.kvs(
+                            TypeDescriptor.of(AvroGeneratedUser.class),
+                            TypeDescriptor.of(TableRow.class))))
+                .via(
+                    kv -> {
+                      String key = kv.getKey();
+                      Iterable<AvroGeneratedUser> il = kv.getValue().getAll(tl);
+                      Iterable<TableRow> ir = kv.getValue().getAll(tr);
+                      List<KV<String, KV<AvroGeneratedUser, TableRow>>> output = new ArrayList<>();
+                      for (AvroGeneratedUser l : il) {
+                        for (TableRow r : ir) {
+                          output.add(KV.of(key, KV.of(l, r)));
+                        }
+                      }
+                      return output;
+                    })
+        ).apply(Count.globally())
+        .apply(
+          MapElements.into(TypeDescriptors.longs())
+              .via(
+                  c -> {
+                    System.out.println("Global count = " + c);
+                    return c;
+                  }));
+
+    final long startTime = System.currentTimeMillis();
+    final State state = pipeline.run().waitUntilFinish();
+    System.out.println(
+        String.format(
+            "CoGroupByKeyBenchmark finished with state %s in %d ms",
+            state, System.currentTimeMillis() - startTime));
+  }
+}