First version of the Iceberg sink for Apache Beam

api/src/main/java/org/apache/iceberg/catalog/Namespace.java

@@ -19,13 +19,14 @@
 
 package org.apache.iceberg.catalog;
 
+import java.io.Serializable;
 import java.util.Arrays;
 import org.apache.iceberg.relocated.com.google.common.base.Joiner;
 
 /**
  * A namespace in a {@link Catalog}.
  */
-public class Namespace {
+public class Namespace implements Serializable {
   private static final Namespace EMPTY_NAMESPACE = new Namespace(new String[] {});
   private static final Joiner DOT = Joiner.on('.');
 

api/src/main/java/org/apache/iceberg/catalog/TableIdentifier.java

@@ -19,6 +19,7 @@
 
 package org.apache.iceberg.catalog;
 
+import java.io.Serializable;
 import java.util.Arrays;
 import java.util.Objects;
 import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
@@ -28,7 +29,7 @@
 /**
  * Identifies a table in iceberg catalog.
  */
-public class TableIdentifier {
+public class TableIdentifier implements Serializable {
 
   private static final Splitter DOT = Splitter.on('.');
 

beam/src/main/java/org/apache/iceberg/beam/IcebergIO.java

@@ -0,0 +1,195 @@
+/*
+ * 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.iceberg.beam;
+
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.transforms.Combine;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.iceberg.AppendFiles;
+import org.apache.iceberg.DataFiles;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.Snapshot;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.catalog.TableIdentifier;
+import org.apache.iceberg.exceptions.NoSuchTableException;
+import org.apache.iceberg.hive.HiveCatalog;
+import org.joda.time.format.DateTimeFormat;
+import org.joda.time.format.DateTimeFormatter;
+
+import java.io.Serializable;
+import java.text.DecimalFormat;
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+
+
+public class IcebergIO {
+
+  private IcebergIO() {
+  }
+
+  static PCollection<Snapshot> write(TableIdentifier table,
+                                     Schema schema,
+                                     String hiveMetastoreUrl,
+                                     WriteFilesResult<Void> resultFiles) {
+    PCollection<KV<Void, String>> filenames = resultFiles
+        .getPerDestinationOutputFilenames();
+
+    FileCombiner combiner = new FileCombiner(table, schema, hiveMetastoreUrl);
+
+    Combine.Globally<KV<Void, String>, Snapshot> combined = Combine.globally(combiner).withoutDefaults();
+
+    return filenames.apply(combined);
+  }
+
+  /**
+   * Defines a custom {@link FileIO.Write.FileNaming} which will use the prefix and suffix supplied to create
+   * a name based on the window, pane, number of shards, shard index, and compression. Removes
+   * window when in the {@link GlobalWindow} and pane info when it is the only firing of the pane.
+   */
+  public static class HadoopCompatibleFilenamePolicy implements FileIO.Write.FileNaming, Serializable {
+    private final String suffix;
+
+    public HadoopCompatibleFilenamePolicy(String suffix) {
+      this.suffix = suffix;
+    }
+
+    @Override
+    public String getFilename(BoundedWindow window, PaneInfo pane, int numShards, int shardIndex, Compression compression) {
+      // Replace the HH:mm:ss with HH-mm-ss, so we don't use any semicolon's
+      final String format = "yyyy-MM-dd'T'HH-mm-ss";
+      final DateTimeFormatter formatter = DateTimeFormat.forPattern(format);
+
+      final StringBuilder res = new StringBuilder();
+      if (window instanceof IntervalWindow) {
+        if (res.length() > 0) {
+          res.append("-");
+        }
+        final IntervalWindow iw = (IntervalWindow) window;
+        res.append(iw.start().toString(formatter)).append("-").append(iw.end().toString(formatter));
+      }
+      boolean isOnlyFiring = pane.isFirst() && pane.isLast();
+      if (!isOnlyFiring) {
+        if (res.length() > 0) {
+          res.append("-");
+        }
+        res.append(pane.getIndex());
+      }
+      if (res.length() > 0) {
+        res.append("-");
+      }
+      final String numShardsStr = String.valueOf(numShards);
+      // A trillion shards per window per pane ought to be enough for everybody.
+      final DecimalFormat df =
+          new DecimalFormat("000000000000".substring(0, Math.max(5, numShardsStr.length())));
+      res.append(df.format(shardIndex)).append("-of-").append(df.format(numShards));
+      res.append(this.suffix);
+      res.append(compression.getSuggestedSuffix());
+      return res.toString();
+    }
+  }
+
+
+  private static class FileCombiner extends Combine.CombineFn<KV<Void, String>, List<String>, Snapshot> {
+    private final TableIdentifier tableIdentifier;
+    private final Schema schema;
+    private final String hiveMetastoreUrl;
+
+    FileCombiner(TableIdentifier table, Schema schema, String hiveMetastoreUrl) {
+      this.tableIdentifier = table;
+      this.schema = schema;
+      this.hiveMetastoreUrl = hiveMetastoreUrl;
+    }
+
+    @Override
+    public List<String> createAccumulator() {
+      return new ArrayList<>();
+    }
+
+    @Override
+    public List<String> addInput(List<String> mutableAccumulator, KV<Void, String> input) {
+      mutableAccumulator.add(input.getValue());
+      return mutableAccumulator;
+    }
+
+    @Override
+    public List<String> mergeAccumulators(Iterable<List<String>> accumulators) {
+      Iterator<List<String>> itr = accumulators.iterator();
+
+      if(itr.hasNext()) {
+        List<String> first = itr.next();
+
+        while(itr.hasNext()) {
+          first.addAll(itr.next());
+        }
+
+        return first;
+      } else {
+        return new ArrayList<>();
+      }
+    }
+
+    @Override
+    public Snapshot extractOutput(List<String> accumulator) {
+      try (HiveCatalog catalog = new HiveCatalog(
+          HiveCatalog.DEFAULT_NAME,
+          this.hiveMetastoreUrl,
+          1,
+          new Configuration()
+      )) {
+        Table table;
+        try {
+          table = catalog.loadTable(this.tableIdentifier);
+        } catch (NoSuchTableException e) {
+          // If it doesn't exist, we just create the table
+          table = catalog.createTable(this.tableIdentifier, schema);
+        }
+
+        // In case the schema has changed
+        if (table.schema() != this.schema) {
+          table.updateSchema().unionByNameWith(this.schema).commit();
+        }
+
+        // Append the new files
+        final AppendFiles app = table.newAppend();
+        // We need to get the statistics, not easy to get them through Beam
+        for (String filename : accumulator) {
+          app.appendFile(DataFiles.builder(table.spec())
+              .withPath(filename)
+              .withFileSizeInBytes(1024)
+              .withRecordCount(100)
+              .build());
+        }
+        app.commit();
+
+        return table.currentSnapshot();
+      }
+    }
+  }
+
+}

beam/src/test/java/org/apache/iceberg/beam/IcebergIOTest.java

@@ -0,0 +1,166 @@
+/*
+ * 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.iceberg.beam;
+
+import org.apache.avro.Schema;
+import org.apache.avro.generic.GenericData;
+import org.apache.avro.generic.GenericRecord;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.coders.AvroCoder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.io.AvroIO;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.testing.TestStream;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.windowing.FixedWindows;
+import org.apache.beam.sdk.transforms.windowing.Window;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.TimestampedValue;
+import org.apache.iceberg.Snapshot;
+import org.apache.iceberg.avro.AvroSchemaUtil;
+import org.apache.iceberg.catalog.TableIdentifier;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.junit.Rule;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.List;
+
+public class IcebergIOTest {
+  private static final List<String> SENTENCES =
+      Arrays.asList(
+          "Beam window 1 1",
+          "Beam window 1 2",
+          "Beam window 1 3",
+          "Beam window 1 4",
+          "Beam window 2 1",
+          "Beam window 2 2");
+  private static final Instant START_TIME = new Instant(0);
+  private static final Duration WINDOW_DURATION = Duration.standardMinutes(1);
+
+  @Rule
+  public final transient TestPipeline pipeline = TestPipeline.create();
+  final String hiveMetastoreUrl = "thrift://localhost:9083/default";
+
+  private static final PipelineOptions options = TestPipeline.testingPipelineOptions();
+
+  private static final String stringSchema = "{\n" +
+      "\t\"type\": \"record\",\n" +
+      "\t\"name\": \"Word\",\n" +
+      "\t\"fields\": [{\n" +
+      "\t\t\"name\": \"word\",\n" +
+      "\t\t\"type\": [\"null\", \"string\"],\n" +
+      "\t\t\"default\": null\n" +
+      "\t}]\n" +
+      "}";
+
+  final Schema avroSchema = new Schema.Parser().parse(stringSchema);
+
+  public static class StringToGenericRecord extends DoFn<String, GenericRecord> {
+    private final Schema schema;
+
+    public StringToGenericRecord() {
+      schema = new Schema.Parser().parse(stringSchema);
+    }
+
+    @ProcessElement
+    public void processElement(@Element String word, OutputReceiver<GenericRecord> out) {
+      GenericRecord record = new GenericData.Record(schema);
+      record.put("word", word);
+      out.output(record);
+    }
+  }
+
+  @Test
+  public void testWriteFilesBatch() {
+    final PipelineOptions options = PipelineOptionsFactory.create();
+    final Pipeline p = Pipeline.create(options);
+
+    p.getCoderRegistry().registerCoderForClass(GenericRecord.class, AvroCoder.of(avroSchema));
+
+    PCollection<String> lines = p.apply(Create.of(SENTENCES)).setCoder(StringUtf8Coder.of());
+
+    PCollection<GenericRecord> records = lines.apply(ParDo.of(new StringToGenericRecord()));
+
+    final String hiveMetastoreUrl = "thrift://localhost:9083/default";
+    FileIO.Write<Void, GenericRecord> avroFileIO = FileIO.<GenericRecord>write()
+        .via(AvroIO.sink(avroSchema))
+        .to("/tmp/fokko/")
+        .withSuffix(".avro");
+
+    WriteFilesResult<Void> filesWritten = records.apply(avroFileIO);
+    org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(avroSchema);
+    TableIdentifier name = TableIdentifier.of("default", "test_batch");
+
+    IcebergIO.write(name, icebergSchema, hiveMetastoreUrl, filesWritten);
+
+    p.run();
+  }
+
+  @Test
+  public void testWriteFilesStreaming() {
+    pipeline.getCoderRegistry().registerCoderForClass(GenericRecord.class, AvroCoder.of(avroSchema));
+
+    // We should see four commits in the log
+    TestStream<String> stringsStream =
+        TestStream.create(StringUtf8Coder.of())
+            .advanceWatermarkTo(START_TIME)
+            .addElements(event(SENTENCES.get(0), 2L))
+            .advanceWatermarkTo(START_TIME.plus(Duration.standardSeconds(60L)))
+            .addElements(event(SENTENCES.get(1), 62L))
+            .advanceWatermarkTo(START_TIME.plus(Duration.standardSeconds(120L)))
+            .addElements(event(SENTENCES.get(2), 122L))
+            .advanceWatermarkTo(START_TIME.plus(Duration.standardSeconds(180L)))
+            .addElements(event(SENTENCES.get(3), 182L))
+            .advanceWatermarkToInfinity();
+
+    PCollection<GenericRecord> recordsStream = pipeline
+        .apply(stringsStream)
+        .setCoder(StringUtf8Coder.of())
+        .apply(ParDo.of(new StringToGenericRecord()));
+
+    FileIO.Write.FileNaming naming = new IcebergIO.HadoopCompatibleFilenamePolicy(".avro");
+
+    FileIO.Write<Void, GenericRecord> avroFileIO = FileIO.<GenericRecord>write()
+        .via(AvroIO.sink(avroSchema))
+        .to("/tmp/fokko/")
+        .withNumShards(1)
+        .withNaming(naming);
+
+    // Write the record
+    WriteFilesResult<Void> filesWritten = recordsStream
+        .apply(Window.into(FixedWindows.of(WINDOW_DURATION)))
+        .apply(avroFileIO);
+
+    org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(avroSchema);
+    TableIdentifier name = TableIdentifier.of("default", "test_streaming5");
+
+    PCollection<Snapshot> snapshots = IcebergIO.write(name, icebergSchema, hiveMetastoreUrl, filesWritten);
+
+    //PAssert.that(snapshots).containsInAnyOrder(new ArrayList<>());
+    pipeline.run(options).waitUntilFinish();
+  }
+
+  private TimestampedValue<String> event(String word, Long timestamp) {
+    return TimestampedValue.of(word, START_TIME.plus(new Duration(timestamp)));
+  }
+}