From e276753290dad4cefcb289b925d12d7287fbec16 Mon Sep 17 00:00:00 2001
From: Anton Okolnychyi <aokolnychyi@apple.com>
Date: Wed, 6 Dec 2023 13:20:10 -0800
Subject: [PATCH] Core: Fix equality in StructLikeMap (#9236)

---
 .../java/org/apache/iceberg/TestHelpers.java  | 46 +++++++++++++++++++
 .../apache/iceberg/util/StructLikeMap.java    | 21 +++------
 .../iceberg/util/TestStructLikeMap.java       | 38 +++++++++++++++
 3 files changed, 91 insertions(+), 14 deletions(-)

diff --git a/api/src/test/java/org/apache/iceberg/TestHelpers.java b/api/src/test/java/org/apache/iceberg/TestHelpers.java
index cefd765a91b5..fcb528caba96 100644
--- a/api/src/test/java/org/apache/iceberg/TestHelpers.java
+++ b/api/src/test/java/org/apache/iceberg/TestHelpers.java
@@ -360,6 +360,52 @@ public int hashCode() {
     }
   }
 
+  // similar to Row but has its own hashCode() and equals() implementations
+  // it is useful for testing custom collections that rely on wrappers
+  public static class CustomRow implements StructLike {
+    public static CustomRow of(Object... values) {
+      return new CustomRow(values);
+    }
+
+    private final Object[] values;
+
+    private CustomRow(Object... values) {
+      this.values = values;
+    }
+
+    @Override
+    public int size() {
+      return values.length;
+    }
+
+    @Override
+    public <T> T get(int pos, Class<T> javaClass) {
+      return javaClass.cast(values[pos]);
+    }
+
+    @Override
+    public <T> void set(int pos, T value) {
+      values[pos] = value;
+    }
+
+    @Override
+    public boolean equals(Object other) {
+      if (this == other) {
+        return true;
+      } else if (other == null || getClass() != other.getClass()) {
+        return false;
+      }
+
+      CustomRow that = (CustomRow) other;
+      return Arrays.equals(values, that.values);
+    }
+
+    @Override
+    public int hashCode() {
+      return 17 * Arrays.hashCode(values);
+    }
+  }
+
   public static class TestFieldSummary implements ManifestFile.PartitionFieldSummary {
     private final boolean containsNull;
     private final Boolean containsNaN;
diff --git a/core/src/main/java/org/apache/iceberg/util/StructLikeMap.java b/core/src/main/java/org/apache/iceberg/util/StructLikeMap.java
index 58bd03041331..2bb5fa1c9d40 100644
--- a/core/src/main/java/org/apache/iceberg/util/StructLikeMap.java
+++ b/core/src/main/java/org/apache/iceberg/util/StructLikeMap.java
@@ -21,7 +21,6 @@
 import java.util.AbstractMap;
 import java.util.Collection;
 import java.util.Map;
-import java.util.Objects;
 import java.util.Set;
 import java.util.function.Function;
 import org.apache.iceberg.StructLike;
@@ -128,9 +127,9 @@ public Set<Entry<StructLike, T>> entrySet() {
 
   private static class StructLikeEntry<R> implements Entry<StructLike, R> {
 
-    private Map.Entry<StructLikeWrapper, R> inner;
+    private final Entry<StructLikeWrapper, R> inner;
 
-    private StructLikeEntry(Map.Entry<StructLikeWrapper, R> inner) {
+    private StructLikeEntry(Entry<StructLikeWrapper, R> inner) {
       this.inner = inner;
     }
 
@@ -146,25 +145,19 @@ public R getValue() {
 
     @Override
     public int hashCode() {
-      int hashCode = getKey().hashCode();
-      if (getValue() != null) {
-        hashCode ^= getValue().hashCode();
-      }
-      return hashCode;
+      return inner.hashCode();
     }
 
     @Override
-    @SuppressWarnings("unchecked")
     public boolean equals(Object o) {
       if (this == o) {
         return true;
-      } else if (!(o instanceof StructLikeEntry)) {
+      } else if (o == null || getClass() != o.getClass()) {
         return false;
-      } else {
-        StructLikeEntry that = (StructLikeEntry<R>) o;
-        return Objects.equals(getKey(), that.getKey())
-            && Objects.equals(getValue(), that.getValue());
       }
+
+      StructLikeEntry<?> that = (StructLikeEntry<?>) o;
+      return inner.equals(that.inner);
     }
 
     @Override
diff --git a/core/src/test/java/org/apache/iceberg/util/TestStructLikeMap.java b/core/src/test/java/org/apache/iceberg/util/TestStructLikeMap.java
index db4176ce231e..f18c48eaa344 100644
--- a/core/src/test/java/org/apache/iceberg/util/TestStructLikeMap.java
+++ b/core/src/test/java/org/apache/iceberg/util/TestStructLikeMap.java
@@ -24,6 +24,8 @@
 import java.util.Map;
 import java.util.Set;
 import org.apache.iceberg.StructLike;
+import org.apache.iceberg.TestHelpers.CustomRow;
+import org.apache.iceberg.TestHelpers.Row;
 import org.apache.iceberg.data.GenericRecord;
 import org.apache.iceberg.data.Record;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
@@ -147,4 +149,40 @@ public void testKeysWithNulls() {
 
     assertThat(map.remove(record3)).isEqualTo("aaa");
   }
+
+  @Test
+  public void testEqualsAndHashCode() {
+    Map<StructLike, String> map1 = StructLikeMap.create(STRUCT_TYPE);
+    Map<StructLike, String> map2 = StructLikeMap.create(STRUCT_TYPE);
+
+    assertThat(map1).isEqualTo(map2);
+    assertThat(map1.hashCode()).isEqualTo(map2.hashCode());
+
+    map1.put(CustomRow.of(1, null), "aaa");
+    map1.put(CustomRow.of(2, null), "bbb");
+
+    map2.put(Row.of(1, null), "aaa");
+    map2.put(Row.of(2, null), "bbb");
+
+    assertThat(map1).isEqualTo(map2);
+    assertThat(map1.hashCode()).isEqualTo(map2.hashCode());
+  }
+
+  @Test
+  public void testKeyAndEntrySetEquality() {
+    Map<StructLike, String> map1 = StructLikeMap.create(STRUCT_TYPE);
+    Map<StructLike, String> map2 = StructLikeMap.create(STRUCT_TYPE);
+
+    assertThat(map1.keySet()).isEqualTo(map2.keySet());
+    assertThat(map1.entrySet()).isEqualTo(map2.entrySet());
+
+    map1.put(CustomRow.of(1, null), "aaa");
+    map1.put(CustomRow.of(2, null), "bbb");
+
+    map2.put(Row.of(1, null), "aaa");
+    map2.put(Row.of(2, null), "bbb");
+
+    assertThat(map1.keySet()).isEqualTo(map2.keySet());
+    assertThat(map1.entrySet()).isEqualTo(map2.entrySet());
+  }
 }