[GitHub] [spark] mridulm commented on a diff in pull request #40121: [SPARK-42528][CORE] Optimize PercentileHeap

["mridulm (via GitHub)" <gi...@apache.org>]

mridulm commented on code in PR #40121:
URL: https://github.com/apache/spark/pull/40121#discussion_r1116742780


##########
core/src/test/scala/org/apache/spark/util/collection/PercentileHeapSuite.scala:
##########
@@ -17,71 +17,73 @@
 
 package org.apache.spark.util.collection
 
-import java.util.NoSuchElementException
+import scala.util.Random
 
 import org.apache.spark.SparkFunSuite
 
 class PercentileHeapSuite extends SparkFunSuite {
 
-  test("If no numbers in PercentileHeap, NoSuchElementException is thrown.") {
-    val medianHeap = new PercentileHeap()
+  test("When PercentileHeap is empty, NoSuchElementException is thrown.") {
+    val medianHeap = new PercentileHeap(0.5)
     intercept[NoSuchElementException] {
       medianHeap.percentile
     }
   }
 
-  test("Median should be correct when size of PercentileHeap is even") {
-    val array = Array(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)
-    val medianHeap = new PercentileHeap()
-    array.foreach(medianHeap.insert(_))
-    assert(medianHeap.size() === 10)
-    assert(medianHeap.smallerSize() === 5)
-    assert(medianHeap.percentile === 4.5)
+  private def percentile(nums: Seq[Int], percentage: Double): Double = {
+    val p = (nums.length * percentage).toInt
+    val sorted = nums.sorted.toIndexedSeq
+    if (nums.length % 2 == 1 || p == 0) {
+      sorted(p)
+    } else {
+      (sorted(p - 1) + sorted(p)) / 2d
+    }
   }
 
-  test("Median should be correct when size of PercentileHeap is odd") {
-    val array = Array(0, 1, 2, 3, 4, 5, 6, 7, 8)
-    val medianHeap = new PercentileHeap()
-    array.foreach(medianHeap.insert(_))
-    assert(medianHeap.size() === 9)
-    assert(medianHeap.smallerSize() === 4)
-    assert(medianHeap.percentile === 4)
+  private def testPercentileFor(nums: Seq[Int], percentage: Double) = {
+    val h = new PercentileHeap(percentage)
+    Random.shuffle(nums).foreach(h.insert(_))
+    assert(h.size == nums.length)
+    assert(h.percentile == percentile(nums, percentage))
   }
 
-  test("Median should be correct though there are duplicated numbers inside.") {
-    val array = Array(0, 0, 1, 1, 2, 3, 4)
-    val medianHeap = new PercentileHeap()
-    array.foreach(medianHeap.insert(_))
-    assert(medianHeap.size === 7)
-    assert(medianHeap.smallerSize() === 3)
-    assert(medianHeap.percentile === 1)
-  }
+  private val tests = Seq(
+    0 until 1,
+    0 until 2,
+    0 until 11,
+    0 until 42,
+    0 until 100
+  )
 
-  test("Median should be correct when input data is skewed.") {
-    val medianHeap = new PercentileHeap()
-    (0 until 10).foreach(_ => medianHeap.insert(5))
-    assert(medianHeap.percentile === 5)
-    (0 until 100).foreach(_ => medianHeap.insert(10))
-    assert(medianHeap.percentile === 10)
-    (0 until 1000).foreach(_ => medianHeap.insert(0))
-    assert(medianHeap.percentile === 0)
+  for (t <- tests) {
+    for (p <- Seq(1, 50, 99)) {
+      test(s"$p% of ${t.mkString(",")}") {
+        testPercentileFor(t, p / 100d)
+      }
+    }
   }
 
-  test("Percentile should be correct when size of PercentileHeap is even") {
-    val array = Array(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)
-    val percentileMap = new PercentileHeap(0.7)
-    array.foreach(percentileMap.insert(_))
-    assert(percentileMap.size() === 10)
-    assert(percentileMap.smallerSize() == 7)
-    assert(percentileMap.percentile === 6.5)
-  }
+  test("benchmark") {

Review Comment:
   The main reason to replace median heap with percentile heap was to generalize it and have it used in other places as well (the median heap impl was quite close to what @alkis has proposed IIRC).
   Given that aspirational goal (even if we are not leveraging it in that fashion right now), might be good to keep it around @cloud-fan  ?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org