add test that writes new data to the branch and reads it

apache · Nov 23, 2023 · 761fdae · 761fdae
1 parent 0127c1d
commit 761fdae
Show file tree

Hide file tree

Showing 3 changed files with 72 additions and 9 deletions.
diff --git a/spark/v3.3/spark/src/test/java/org/apache/iceberg/spark/source/TestSnapshotSelection.java b/spark/v3.3/spark/src/test/java/org/apache/iceberg/spark/source/TestSnapshotSelection.java
@@ -427,11 +427,32 @@ public void testSnapshotSelectionByBranchWithSchemaChange() throws IOException {
                 .option("branch", "branch")
                 .load(tableLocation)
                 .orderBy("id")
+                .as(Encoders.bean(SimpleRecord.class))
                 .collectAsList())
         .containsExactly(
-            new GenericRowWithSchema(new Object[] {1, null}, null),
-            new GenericRowWithSchema(new Object[] {2, null}, null),
-            new GenericRowWithSchema(new Object[] {3, null}, null));
+            new SimpleRecord(1, null), new SimpleRecord(2, null), new SimpleRecord(3, null));
+
+    // writing new records into the branch should work with the re-introduced column
+    List<SimpleRecord> records =
+        Lists.newArrayList(
+            new SimpleRecord(4, "d"), new SimpleRecord(5, "e"), new SimpleRecord(6, "f"));
+    Dataset<Row> dataFrame = spark.createDataFrame(records, SimpleRecord.class);
+    dataFrame
+        .select("id", "data")
+        .write()
+        .format("iceberg")
+        .option("branch", "branch")
+        .mode("append")
+        .save(tableLocation);
+
+    Assertions.assertThat(
+            branchSnapshotResult
+                .orderBy("id")
+                .as(Encoders.bean(SimpleRecord.class))
+                .collectAsList())
+        .hasSize(6)
+        .contains(new SimpleRecord(1, null), new SimpleRecord(2, null), new SimpleRecord(3, null))
+        .containsAll(records);
   }
 
   @Test

diff --git a/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestSnapshotSelection.java b/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestSnapshotSelection.java
@@ -450,11 +450,32 @@ public void testSnapshotSelectionByBranchWithSchemaChange() throws IOException {
                 .option("branch", "branch")
                 .load(tableLocation)
                 .orderBy("id")
+                .as(Encoders.bean(SimpleRecord.class))
                 .collectAsList())
         .containsExactly(
-            new GenericRowWithSchema(new Object[] {1, null}, null),
-            new GenericRowWithSchema(new Object[] {2, null}, null),
-            new GenericRowWithSchema(new Object[] {3, null}, null));
+            new SimpleRecord(1, null), new SimpleRecord(2, null), new SimpleRecord(3, null));
+
+    // writing new records into the branch should work with the re-introduced column
+    List<SimpleRecord> records =
+        Lists.newArrayList(
+            new SimpleRecord(4, "d"), new SimpleRecord(5, "e"), new SimpleRecord(6, "f"));
+    Dataset<Row> dataFrame = spark.createDataFrame(records, SimpleRecord.class);
+    dataFrame
+        .select("id", "data")
+        .write()
+        .format("iceberg")
+        .option("branch", "branch")
+        .mode("append")
+        .save(tableLocation);
+
+    Assertions.assertThat(
+            branchSnapshotResult
+                .orderBy("id")
+                .as(Encoders.bean(SimpleRecord.class))
+                .collectAsList())
+        .hasSize(6)
+        .contains(new SimpleRecord(1, null), new SimpleRecord(2, null), new SimpleRecord(3, null))
+        .containsAll(records);
   }
 
   @Test

diff --git a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/source/TestSnapshotSelection.java b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/source/TestSnapshotSelection.java
@@ -450,11 +450,32 @@ public void testSnapshotSelectionByBranchWithSchemaChange() throws IOException {
                 .option("branch", "branch")
                 .load(tableLocation)
                 .orderBy("id")
+                .as(Encoders.bean(SimpleRecord.class))
                 .collectAsList())
         .containsExactly(
-            new GenericRowWithSchema(new Object[] {1, null}, null),
-            new GenericRowWithSchema(new Object[] {2, null}, null),
-            new GenericRowWithSchema(new Object[] {3, null}, null));
+            new SimpleRecord(1, null), new SimpleRecord(2, null), new SimpleRecord(3, null));
+
+    // writing new records into the branch should work with the re-introduced column
+    List<SimpleRecord> records =
+        Lists.newArrayList(
+            new SimpleRecord(4, "d"), new SimpleRecord(5, "e"), new SimpleRecord(6, "f"));
+    Dataset<Row> dataFrame = spark.createDataFrame(records, SimpleRecord.class);
+    dataFrame
+        .select("id", "data")
+        .write()
+        .format("iceberg")
+        .option("branch", "branch")
+        .mode("append")
+        .save(tableLocation);
+
+    Assertions.assertThat(
+            branchSnapshotResult
+                .orderBy("id")
+                .as(Encoders.bean(SimpleRecord.class))
+                .collectAsList())
+        .hasSize(6)
+        .contains(new SimpleRecord(1, null), new SimpleRecord(2, null), new SimpleRecord(3, null))
+        .containsAll(records);
   }
 
   @Test