Relax ordering validation

spotify · Nov 21, 2024 · e334184 · e334184
1 parent 34378a8
commit e334184
Show file tree

Hide file tree

Showing 5 changed files with 7 additions and 47 deletions.
diff --git a/scio-core/src/main/scala/com/spotify/scio/coders/instances/JavaCoders.scala b/scio-core/src/main/scala/com/spotify/scio/coders/instances/JavaCoders.scala
@@ -72,13 +72,8 @@ final private[coders] class JPriorityQueueCoder[T](
     pq
   }
 
-  override def encode(value: java.util.PriorityQueue[T], os: OutputStream): Unit = {
-    require(
-      value.comparator() == ordering,
-      "PriorityQueue comparator does not match JPriorityQueueCoder comparator"
-    )
+  override def encode(value: java.util.PriorityQueue[T], os: OutputStream): Unit =
     super.encode(value, os)
-  }
 
   override def verifyDeterministic(): Unit =
     throw new NonDeterministicException(

diff --git a/scio-core/src/main/scala/com/spotify/scio/coders/instances/ScalaCoders.scala b/scio-core/src/main/scala/com/spotify/scio/coders/instances/ScalaCoders.scala
@@ -245,13 +245,8 @@ private[coders] class MutableSetCoder[T](bc: BCoder[T]) extends SeqLikeCoder[m.S
 
 private class SortedSetCoder[T: Ordering](bc: BCoder[T]) extends SeqLikeCoder[SortedSet, T](bc) {
 
-  override def encode(value: SortedSet[T], os: OutputStream): Unit = {
-    require(
-      value.ordering == Ordering[T],
-      "SortedSet ordering does not match SortedSetCoder ordering"
-    )
+  override def encode(value: SortedSet[T], os: OutputStream): Unit =
     super.encode(value, os)
-  }
 
   override def decode(inStream: InputStream): SortedSet[T] =
     decode(inStream, SortedSet.newBuilder[T])
@@ -260,13 +255,8 @@ private class SortedSetCoder[T: Ordering](bc: BCoder[T]) extends SeqLikeCoder[So
 private[coders] class MutablePriorityQueueCoder[T: Ordering](bc: BCoder[T])
     extends SeqLikeCoder[m.PriorityQueue, T](bc) {
   override def consistentWithEquals(): Boolean = false // PriorityQueue does not define equality
-  override def encode(value: m.PriorityQueue[T], os: OutputStream): Unit = {
-    require(
-      value.ord == Ordering[T],
-      "PriorityQueue ordering does not match MutablePriorityQueueCoder ordering"
-    )
+  override def encode(value: m.PriorityQueue[T], os: OutputStream): Unit =
     super.encode(value, os)
-  }
   override def decode(inStream: InputStream): m.PriorityQueue[T] =
     decode(inStream, m.PriorityQueue.newBuilder[T])
 }
@@ -377,13 +367,8 @@ private class MutableMapCoder[K, V](kc: BCoder[K], vc: BCoder[V])
 private[coders] class SortedMapCoder[K: Ordering, V](kc: BCoder[K], vc: BCoder[V])
     extends MapLikeCoder[K, V, SortedMap](kc, vc) {
 
-  override def encode(value: SortedMap[K, V], os: OutputStream): Unit = {
-    require(
-      value.ordering == Ordering[K],
-      "SortedMap ordering does not match SortedMapCoder ordering"
-    )
+  override def encode(value: SortedMap[K, V], os: OutputStream): Unit =
     super.encode(value, os)
-  }
 
   override def decode(is: InputStream): SortedMap[K, V] =
     decode(is, SortedMap.newBuilder[K, V])

diff --git a/scio-core/src/main/scala/com/spotify/scio/values/SampleSCollectionFunctions.scala b/scio-core/src/main/scala/com/spotify/scio/values/SampleSCollectionFunctions.scala
@@ -30,13 +30,8 @@ object SampleSCollectionFunctions {
 
   final private case class WeightedSample[T](id: Long, value: T, weight: Long)
   private object WeightedSample {
-    private object WeightedSampleOrdering extends Ordering[WeightedSample[_]] {
-      override def compare(x: WeightedSample[_], y: WeightedSample[_]): Int =
-        java.lang.Long.compare(y.id, x.id) // reverse order
-    }
-
     implicit def ordering[T]: Ordering[WeightedSample[T]] =
-      WeightedSampleOrdering.asInstanceOf[Ordering[WeightedSample[T]]]
+      Ordering.by[WeightedSample[T], Long](_.id).reverse
   }
 
   final private case class WeightedCombiner[T](

diff --git a/scio-core/src/test/scala/com/spotify/scio/coders/CoderTest.scala b/scio-core/src/test/scala/com/spotify/scio/coders/CoderTest.scala
@@ -158,8 +158,7 @@ final class CoderTest extends AnyFlatSpec with Matchers {
       beFullyCompliantNotConsistentWithEquals()
 
     {
-      // custom ordering must have stable equal after serialization
-      implicit val pqOrd: Ordering[String] = FlippedStringOrdering
+      implicit val pqOrd: Ordering[String] = Ordering.String.on(_.reverse)
       val pq = new mut.PriorityQueue[String]()(pqOrd)
       pq ++= s
 
@@ -416,8 +415,7 @@ final class CoderTest extends AnyFlatSpec with Matchers {
     }
 
     {
-      // custom ordering must have stable equal after serialization
-      implicit val pqOrd: Ordering[String] = FlippedStringOrdering
+      implicit val pqOrd: Ordering[String] = Ordering.String.on(_.reverse)
       val pq = new JPriorityQueue[String](pqOrd)
       pq.addAll(elems.asJavaCollection)
 
@@ -995,15 +993,6 @@ final class CoderTest extends AnyFlatSpec with Matchers {
     )
   }
 
-  it should "not accept SortedMap when ordering doesn't match with coder" in {
-    val sm = SortedMap(1 -> "1", 2 -> "2")(Ordering[Int].reverse)
-    // implicit SortedMapCoder will use implicit default Int ordering
-    val e = the[IllegalArgumentException] thrownBy {
-      sm coderShould roundtrip()
-    }
-    e.getMessage shouldBe "requirement failed: SortedMap ordering does not match SortedMapCoder ordering"
-  }
-
   /*
    * Case class nested inside another class. Do not move outside
    * */

diff --git a/scio-core/src/test/scala/com/spotify/scio/coders/CoderTestUtils.scala b/scio-core/src/test/scala/com/spotify/scio/coders/CoderTestUtils.scala
@@ -29,10 +29,6 @@ import scala.jdk.CollectionConverters._
 
 object CoderTestUtils {
 
-  object FlippedStringOrdering extends Ordering[String] {
-    override def compare(x: String, y: String): Int = x.reverse.compareTo(y.reverse)
-  }
-
   def testRoundTrip[T](coder: BCoder[T], value: T): Boolean =
     testRoundTrip(coder, coder, value)