StreamingOffset Of Structured streaming read for Iceberg

spark/src/main/java/org/apache/iceberg/spark/source/StreamingOffset.java

@@ -0,0 +1,150 @@
+/*
+ * 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.spark.source;
+
+import com.fasterxml.jackson.core.JsonGenerator;
+import com.fasterxml.jackson.databind.JsonNode;
+import java.io.IOException;
+import java.io.StringWriter;
+import java.io.UncheckedIOException;
+import org.apache.iceberg.relocated.com.google.common.base.Objects;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.util.JsonUtil;
+import org.apache.spark.sql.sources.v2.reader.streaming.Offset;
+
+/**
+ * An implementation of Spark Structured Streaming Offset, to track the current processed
+ * files of Iceberg table. This StreamingOffset consists of:
+ *
+ * version: The version of StreamingOffset. The offset was created with a version number used to validate
+ * when deserializing from json string.
+ * snapshot_id: The current processed snapshot id.
+ * index: The index of last scanned file in snapshot.
+ * scan_all_files: Denote whether to scan all files in a snapshot, currently we only scan all files in the starting
+ * snapshot.
+ * snapshot_fully_processed: Denote whether the current snapshot is fully processed, to avoid revisiting the processed
+ * snapshot.
+ */
+class StreamingOffset extends Offset {
+  static final StreamingOffset START_OFFSET = new StreamingOffset(-1L, -1, false, true);
+
+  private static final int CURR_VERSION = 1;
+  private static final String VERSION = "version";
+  private static final String SNAPSHOT_ID = "snapshot_id";
+  private static final String INDEX = "index";
+  private static final String SCAN_ALL_FILES = "scan_all_files";
+  private static final String SNAPSHOT_FULLY_PROCESSED = "snapshot_fully_processed";
+
+  private final long snapshotId;
+  private final int index;
+  private final boolean scanAllFiles;
+  private final boolean snapshotFullyProcessed;
+
+  StreamingOffset(long snapshotId, int index, boolean scanAllFiles, boolean snapshotFullyProcessed) {
+    this.snapshotId = snapshotId;
+    this.index = index;
+    this.scanAllFiles = scanAllFiles;
+    this.snapshotFullyProcessed = snapshotFullyProcessed;
+  }
+
+  static StreamingOffset fromJson(String json) {
+    Preconditions.checkNotNull(json, "The input JSON string is null");
+
+    try {
+      JsonNode node = JsonUtil.mapper().readValue(json, JsonNode.class);
+      int version = JsonUtil.getInt(VERSION, node);
+      if (version > CURR_VERSION) {
+        throw new IOException(String.format("Cannot deserialize a JSON offset from version %d. %d is not compatible " +
+            "with the version of Iceberg %d and cannot be used. Please use a compatible version of Iceberg " +
+            "to read this offset", version, version, CURR_VERSION));
+      }
+
+      long snapshotId = JsonUtil.getLong(SNAPSHOT_ID, node);
+      int index = JsonUtil.getInt(INDEX, node);
+      boolean shouldScanAllFiles = JsonUtil.getBool(SCAN_ALL_FILES, node);
+      boolean snapshotFullyProcessed = JsonUtil.getBool(SNAPSHOT_FULLY_PROCESSED, node);
+
+      return new StreamingOffset(snapshotId, index, shouldScanAllFiles, snapshotFullyProcessed);
+    } catch (IOException e) {
+      throw new IllegalStateException(String.format("Failed to parse StreamingOffset from JSON string %s", json), e);
+    }
+  }
+
+  @Override
+  public String json() {
+    StringWriter writer = new StringWriter();
+    try {
+      JsonGenerator generator = JsonUtil.factory().createGenerator(writer);
+      generator.writeStartObject();
+      generator.writeNumberField(VERSION, CURR_VERSION);
+      generator.writeNumberField(SNAPSHOT_ID, snapshotId);
+      generator.writeNumberField(INDEX, index);
+      generator.writeBooleanField(SCAN_ALL_FILES, scanAllFiles);
+      generator.writeBooleanField(SNAPSHOT_FULLY_PROCESSED, snapshotFullyProcessed);
+      generator.writeEndObject();
+      generator.flush();
+
+    } catch (IOException e) {
+      throw new UncheckedIOException("Failed to write StreamingOffset to json", e);
+    }
+
+    return writer.toString();
+  }
+
+  long snapshotId() {
+    return snapshotId;
+  }
+
+  int index() {
+    return index;
+  }
+
+  boolean shouldScanAllFiles() {
+    return scanAllFiles;
+  }
+
+  boolean isSnapshotFullyProcessed() {
+    return snapshotFullyProcessed;
+  }
+
+  @Override
+  public boolean equals(Object obj) {
+    if (obj instanceof StreamingOffset) {
+      StreamingOffset offset = (StreamingOffset) obj;
+      return offset.snapshotId == snapshotId &&
+          offset.index == index &&
+          offset.scanAllFiles == scanAllFiles &&
+          offset.snapshotFullyProcessed == snapshotFullyProcessed;
+    } else {
+      return false;
+    }
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hashCode(snapshotId, index, scanAllFiles, snapshotFullyProcessed);
+  }
+
+  @Override
+  public String toString() {
+    return String.format("Streaming Offset[%d: index (%d) scan_all_files (%b) snapshot_fully_processed (%b)]",
+      snapshotId, index, scanAllFiles, snapshotFullyProcessed);
+  }
+}

spark/src/test/java/org/apache/iceberg/spark/source/TestStreamingOffset.java

@@ -0,0 +1,105 @@
+/*
+ * 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.spark.source;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.apache.iceberg.DataFile;
+import org.apache.iceberg.ManifestFile;
+import org.apache.iceberg.ManifestFiles;
+import org.apache.iceberg.ManifestReader;
+import org.apache.iceberg.Snapshot;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.junit.Assert;
+import org.junit.Test;
+
+public abstract class TestStreamingOffset extends TestStructuredStreamingRead {
+
+  @Test
+  public void testStreamingOffsetWithPosition() throws IOException {
+    Snapshot currSnap = table.currentSnapshot();
+    StreamingOffset startOffset =
+        new StreamingOffset(currSnap.snapshotId(), INIT_SCANNED_FILE_INDEX, false, true);
+    StreamingOffset endOffset = startOffset;
+    ManifestFile manifest = currSnap.dataManifests().get(0);
+    try (ManifestReader<DataFile> reader = ManifestFiles.read(manifest, FILE_IO)) {
+      long expectedPos = startOffset.index();
+      for (DataFile file : reader) {
+        expectedPos += 1;
+        Assert.assertEquals("Position should match", (Long) expectedPos, file.pos());
+        endOffset =
+            new StreamingOffset(currSnap.snapshotId(), Math.toIntExact(file.pos()), false, true);
+      }
+      StreamingOffset expectedOffset =
+          new StreamingOffset(currSnap.snapshotId(), (int) expectedPos, false, true);
+      Assert.assertEquals(expectedOffset, endOffset);
+    }
+  }
+
+  @Test
+  public void testScanAllFiles() throws IOException {
+    List<ManifestFile> manifests = table.currentSnapshot().dataManifests();
+    List<StreamingOffset> expectedOffsets = Lists.newArrayList();
+    for (ManifestFile manifest : manifests) {
+      expectedOffsets
+          .add(new StreamingOffset(manifest.snapshotId(), END_SCANNED_FILE_INDEX, true, true));
+    }
+    testStreamingOffsetWithScanFiles(expectedOffsets, true);
+  }
+
+  @Test
+  public void testNoScanAllFiles() throws IOException {
+    List<StreamingOffset> expectedOffsets = Lists
+        .newArrayList(
+            new StreamingOffset(table.currentSnapshot().snapshotId(), END_SCANNED_FILE_INDEX, false,
+                true));
+    testStreamingOffsetWithScanFiles(expectedOffsets, false);
+  }
+
+  private void testStreamingOffsetWithScanFiles(List<StreamingOffset> expectedOffsets,
+      boolean scanAllFiles) throws IOException {
+    Snapshot currSnap = table.currentSnapshot();
+    List<ManifestFile> manifests = scanAllFiles ? currSnap.dataManifests() :
+        currSnap.dataManifests().stream().filter(m -> m.snapshotId().equals(currSnap.snapshotId()))
+            .collect(Collectors.toList());
+    List<StreamingOffset> actualOffsets = Lists.newArrayList();
+    for (ManifestFile manifest : manifests) {
+      try (ManifestReader<DataFile> reader = ManifestFiles.read(manifest, FILE_IO)) {
+        StreamingOffset offset = StreamingOffset.START_OFFSET;
+        long expectedPos = INIT_SCANNED_FILE_INDEX;
+        for (DataFile file : reader) {
+          expectedPos += 1;
+          Assert.assertEquals("Position should match", (Long) expectedPos, file.pos());
+          if (scanAllFiles) {
+            offset = new StreamingOffset(manifest.snapshotId(), Math.toIntExact(file.pos()),
+                scanAllFiles, true);
+          } else {
+            offset = new StreamingOffset(currSnap.snapshotId(), Math.toIntExact(file.pos()),
+                scanAllFiles, true);
+          }
+        }
+        actualOffsets.add(offset);
+      }
+    }
+
+    Assert.assertArrayEquals(expectedOffsets.toArray(), actualOffsets.toArray());
+  }
+}

spark/src/test/java/org/apache/iceberg/spark/source/TestStructuredStreamingRead.java

@@ -0,0 +1,104 @@
+/*
+ * 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.spark.source;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.hadoop.HadoopTables;
+import org.apache.iceberg.io.FileIO;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.types.Types;
+import org.apache.spark.sql.Dataset;
+import org.apache.spark.sql.Row;
+import org.apache.spark.sql.SparkSession;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.rules.ExpectedException;
+import org.junit.rules.TemporaryFolder;
+
+import static org.apache.iceberg.types.Types.NestedField.optional;
+
+public abstract class TestStructuredStreamingRead {
+
+  private static final Configuration CONF = new Configuration();
+  private static final Schema SCHEMA = new Schema(
+      optional(1, "id", Types.IntegerType.get()),
+      optional(2, "data", Types.StringType.get())
+  );
+  protected static final int INIT_SCANNED_FILE_INDEX = -1;
+  protected static final int END_SCANNED_FILE_INDEX = 3;
+  protected static SparkSession spark = null;
+  protected static Path parent = null;
+  protected static File tableLocation = null;
+  protected static Table table = null;
+  protected static List<SimpleRecord> expected = null;
+  protected static final FileIO FILE_IO = new TestTables.LocalFileIO();
+
+  @Rule
+  public TemporaryFolder temp = new TemporaryFolder();
+  @Rule
+  public ExpectedException exceptionRule = ExpectedException.none();
+
+  @BeforeClass
+  public static void startSpark() throws IOException {
+    TestStructuredStreamingRead.spark = SparkSession.builder()
+        .master("local[2]")
+        .config("spark.sql.shuffle.partitions", 4)
+        .getOrCreate();
+
+    parent = Files.createTempDirectory("test");
+    tableLocation = new File(parent.toFile(), "table");
+    tableLocation.mkdir();
+    HadoopTables tables = new HadoopTables(CONF);
+    PartitionSpec spec = PartitionSpec.builderFor(SCHEMA).identity("data").build();
+    table = tables.create(SCHEMA, spec, tableLocation.toString());
+
+    expected = Lists.newArrayList(new SimpleRecord(1, "1"),
+        new SimpleRecord(2, "2"),
+        new SimpleRecord(3, "3"),
+        new SimpleRecord(4, "4"));
+
+    // Write records one by one to generate 3 snapshots.
+    for (int i = 0; i < 3; i++) {
+      Dataset<Row> df = spark.createDataFrame(expected, SimpleRecord.class);
+      df.select("id", "data").write()
+          .format("iceberg")
+          .mode("append")
+          .save(tableLocation.toString());
+    }
+
+    table.refresh();
+  }
+
+  @AfterClass
+  public static void stopSpark() {
+    SparkSession currentSpark = TestStructuredStreamingRead.spark;
+    TestStructuredStreamingRead.spark = null;
+    currentSpark.stop();
+  }
+}

spark2/src/test/java/org/apache/iceberg/spark/source/TestStreamingOffset24.java

@@ -0,0 +1,24 @@
+/*
+ * 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.spark.source;
+
+public class TestStreamingOffset24 extends TestStreamingOffset {
+
+}