Repository: G-Research/spark-extension (language: python)

Issue:
MapComparator is sensitive to key order
As pointed out in https://github.com/G-Research/spark-extension/issues/184#issuecomment-1689594998, this example shows that the map comparator identifies a `C`hange where it should flag `N`o change:

```scala
val left = Seq((1, """{"foo":"__foo__","bar":"__bar__"}""")).toDF("id", "json")
val right = Seq((1, """{"bar":"__bar__","foo":"__foo__"}""")).toDF("id", "json")
val leftMap = left.select($"id", from_json($"json", MapType(StringType, StringType)))
val rightMap = right.select($"id", from_json($"json", MapType(StringType, StringType)))
val options = DiffOptions.default.withComparator(DiffComparators.map[String, String], MapType(StringType, StringType))
new Differ(options).diff(leftMap, rightMap, "id").show
```
```
+----+---+--------------------------------+--------------------------------+
|diff|id |left_entries                    |right_entries                   |
+----+---+--------------------------------+--------------------------------+
|C   |1  |{foo -> __foo__, bar -> __bar__}|{bar -> __bar__, foo -> __foo__}|
+----+---+--------------------------------+--------------------------------+
```


Discussion / hints:

Fix patch:
diff --git a/src/main/scala/uk/co/gresearch/spark/diff/comparator/MapDiffComparator.scala b/src/main/scala/uk/co/gresearch/spark/diff/comparator/MapDiffComparator.scala
index 05710f25..b8be8532 100644
--- a/src/main/scala/uk/co/gresearch/spark/diff/comparator/MapDiffComparator.scala
+++ b/src/main/scala/uk/co/gresearch/spark/diff/comparator/MapDiffComparator.scala
@@ -21,32 +21,32 @@ import org.apache.spark.sql.catalyst.expressions.UnsafeMapData
 import org.apache.spark.sql.types.{DataType, MapType}
 import org.apache.spark.sql.{Column, Encoder}
 
+import scala.reflect.ClassTag
+
 case class MapDiffComparator[K, V](private val comparator: EquivDiffComparator[UnsafeMapData]) extends DiffComparator {
   override def equiv(left: Column, right: Column): Column = comparator.equiv(left, right)
 }
 
-private case class MapDiffEquiv[K, V](keyType: DataType, valueType: DataType) extends math.Equiv[UnsafeMapData] {
+private case class MapDiffEquiv[K: ClassTag, V](keyType: DataType, valueType: DataType) extends math.Equiv[UnsafeMapData] {
   override def equiv(left: UnsafeMapData, right: UnsafeMapData): Boolean = {
-    val leftKeys: Map[K, Int] = 0.until(left.keyArray().numElements())
-      .map(ordinal => left.keyArray().get(ordinal, keyType).asInstanceOf[K] -> ordinal)
-      .toMap
-    val rightKeys: Map[K, Int] = 0.until(left.keyArray().numElements())
-      .map(ordinal => left.keyArray().get(ordinal, keyType).asInstanceOf[K] -> ordinal)
-      .toMap
+
+    val leftKeysIndices: Map[K, Int] = left.keyArray().toArray(keyType).zipWithIndex.toMap
+    val rightKeysIndices: Map[K, Int] = right.keyArray().toArray(keyType).zipWithIndex.toMap
 
     val leftValues = left.valueArray()
     val rightValues = right.valueArray()
 
-    val valuesAreEqual = leftKeys
-      .map { case (key, ordinal) => ordinal -> rightKeys(key) }
-      .map { case (leftOrdinal, rightOrdinal) => (leftOrdinal, rightOrdinal, leftValues.isNullAt(leftOrdinal), rightValues.isNullAt(rightOrdinal)) }
-      .map { case (leftOrdinal, rightOrdinal, leftIsNull, rightIsNull) =>
+    // can only be evaluated when right has same keys as left
+    lazy val valuesAreEqual = leftKeysIndices
+      .map { case (key, index) => index -> rightKeysIndices(key) }
+      .map { case (leftIndex, rightIndex) => (leftIndex, rightIndex, leftValues.isNullAt(leftIndex), rightValues.isNullAt(rightIndex)) }
+      .map { case (leftIndex, rightIndex, leftIsNull, rightIsNull) =>
         leftIsNull && rightIsNull ||
-          !leftIsNull && !rightIsNull && leftValues.get(leftOrdinal, valueType).equals(rightValues.get(rightOrdinal, valueType))
+          !leftIsNull && !rightIsNull && leftValues.get(leftIndex, valueType).equals(rightValues.get(rightIndex, valueType))
       }
 
     left.numElements() == right.numElements() &&
-      leftKeys.keySet.diff(rightKeys.keySet).isEmpty &&
+      leftKeysIndices.keySet.diff(rightKeysIndices.keySet).isEmpty &&
       valuesAreEqual.forall(identity)
   }
 }


Test patch:
diff --git a/src/test/scala/uk/co/gresearch/spark/diff/DiffComparatorSuite.scala b/src/test/scala/uk/co/gresearch/spark/diff/DiffComparatorSuite.scala
index 2285cf50..d087dc74 100644
--- a/src/test/scala/uk/co/gresearch/spark/diff/DiffComparatorSuite.scala
+++ b/src/test/scala/uk/co/gresearch/spark/diff/DiffComparatorSuite.scala
@@ -110,6 +110,8 @@ class DiffComparatorSuite extends AnyFunSuite with SparkTestSession {
     Maps(2, Map(1 -> 2L, 2 -> 2L, 3 -> 3L)),
     Maps(3, Map(1 -> 3L, 2 -> 2L, 3 -> 3L)),
     Maps(4, Map(1 -> 4L, 2 -> 2L, 3 -> 3L)),
+    Maps(6, Map(1 -> 1L, 2 -> 2L, 3 -> 3L)),
+    Maps(7, Map(1 -> 4L, 2 -> 2L, 3 -> 3L)),
   ).toDS()
 
   lazy val rightMaps: Dataset[Maps] = Seq(
@@ -117,6 +119,8 @@ class DiffComparatorSuite extends AnyFunSuite with SparkTestSession {
     Maps(2, Map(1 -> 2L, 2 -> 3L, 3 -> 3L)),
     Maps(3, Map(1 -> 3L, 2 -> 2L, 4 -> 4L)),
     Maps(5, Map(1 -> 4L, 2 -> 2L, 3 -> 3L)),
+    Maps(6, Map(3 -> 3L, 2 -> 2L, 1 -> 1L)),
+    Maps(7, Map(3 -> 4L, 2 -> 2L, 1 -> 1L)),
   ).toDS()
 
   def doTest(optionsWithTightComparators: DiffOptions,
@@ -360,7 +364,7 @@ class DiffComparatorSuite extends AnyFunSuite with SparkTestSession {
         val options = DiffOptions.default.withComparator(DiffComparators.map[Int, Long](), "map")
 
         val actual = leftMaps.diff(rightMaps, options, "id").orderBy($"id").collect()
-        val diffs = Seq((1, "N"), (2, "C"), (3, "C"), (4, "D"), (5, "I")).toDF("id", "diff")
+        val diffs = Seq((1, "N"), (2, "C"), (3, "C"), (4, "D"), (5, "I"), (6, "N"), (7, "C")).toDF("id", "diff")
         val expected = leftMaps.withColumnRenamed("map", "left_map")
           .join(rightMaps.withColumnRenamed("map", "right_map"), Seq("id"), "fullouter")
           .join(diffs, "id")