wrap

Author: turboFei

amoro-ams/src/test/java/org/apache/amoro/server/optimizing/flow/view/KeyedTableDataView.java

@@ -34,6 +34,7 @@
 import org.apache.iceberg.PartitionField;
 import org.apache.iceberg.RowDelta;
 import org.apache.iceberg.Schema;
+import org.apache.iceberg.data.InternalRecordWrapper;
 import org.apache.iceberg.data.Record;
 import org.apache.iceberg.io.WriteResult;
 import org.apache.iceberg.types.Type;
@@ -65,6 +66,8 @@ public class KeyedTableDataView extends AbstractTableDataView {
 
   private final RandomRecordGenerator generator;
 
+  private final InternalRecordWrapper wrapper;
+
   public KeyedTableDataView(
       MixedTable mixedTable,
       Schema primary,
@@ -77,6 +80,7 @@ public KeyedTableDataView(
     org.apache.amoro.shade.guava32.com.google.common.base.Preconditions.checkArgument(
         primary.columns().size() == 1
             && primary.columns().get(0).type().typeId() == Type.TypeID.INTEGER);
+    this.wrapper = new InternalRecordWrapper(schema.asStruct());
     this.schemaSize = schema.columns().size();
 
     this.primaryUpperBound = primaryUpperBound;
@@ -87,7 +91,7 @@ public KeyedTableDataView(
     this.view = StructLikeMap.create(primary.asStruct());
     List<Record> records = new DataReader(mixedTable).allData();
     for (Record record : records) {
-      view.put(record, record);
+      view.put(wrapper.copyFor(record), record);
     }
 
     Map<Integer, Map<Integer, Object>> primaryRelationWithPartition = new HashMap<>();
@@ -132,7 +136,7 @@ public WriteResult append(int count) throws IOException {
     List<RecordWithAction> records = new ArrayList<>();
     for (int i = 0; i < primaryUpperBound; i++) {
       Record record = generator.randomRecord(i);
-      if (!view.containsKey(record)) {
+      if (!view.containsKey(wrapper.copyFor(record))) {
         records.add(new RecordWithAction(record, ChangeAction.INSERT));
       }
       if (records.size() == count) {
@@ -156,13 +160,14 @@ public WriteResult cdc(int count) throws IOException {
     List<Record> scatter = randomRecord(count);
     List<RecordWithAction> cdc = new ArrayList<>();
     for (Record record : scatter) {
-      if (view.containsKey(record)) {
+      if (view.containsKey(wrapper.copyFor(record))) {
         if (random.nextBoolean()) {
           // delete
-          cdc.add(new RecordWithAction(view.get(record), ChangeAction.DELETE));
+          cdc.add(new RecordWithAction(view.get(wrapper.copyFor(record)), ChangeAction.DELETE));
         } else {
           // update
-          cdc.add(new RecordWithAction(view.get(record), ChangeAction.UPDATE_BEFORE));
+          cdc.add(
+              new RecordWithAction(view.get(wrapper.copyFor(record)), ChangeAction.UPDATE_BEFORE));
           cdc.add(new RecordWithAction(record, ChangeAction.UPDATE_AFTER));
         }
       } else {
@@ -210,7 +215,7 @@ public MatchResult match(List<Record> records) {
     List<Record> inViewButDuplicate = new ArrayList<>();
     StructLikeSet intersection = StructLikeSet.create(schema.asStruct());
     for (Record record : records) {
-      Record viewRecord = view.get(record);
+      Record viewRecord = view.get(wrapper.copyFor(record));
       if (viewRecord == null) {
         notInView.add(record);
       }
@@ -304,12 +309,12 @@ private void writeView(List<RecordWithAction> records) {
       changeLog.add(record);
       ChangeAction action = record.getAction();
       if (action == ChangeAction.DELETE || action == ChangeAction.UPDATE_BEFORE) {
-        view.remove(record);
+        view.remove(wrapper.copyFor(record));
       } else {
-        if (view.containsKey(record)) {
+        if (view.containsKey(wrapper.copyFor(record))) {
           throw new RuntimeException("You write duplicate pk");
         }
-        view.put(record, record);
+        view.put(wrapper.copyFor(record), record);
       }
     }
   }