Add ordering test

Qbeast-io · Dec 16, 2024 · 57bb975 · 57bb975
1 parent ef3928f
commit 57bb975
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 7 deletions.
diff --git a/core/src/main/scala/io/qbeast/spark/writer/Rollup.scala b/core/src/main/scala/io/qbeast/spark/writer/Rollup.scala
@@ -54,7 +54,7 @@ private[writer] class Rollup(limit: Double) {
    *   the rollup result
    */
   def compute(): Map[CubeId, CubeId] = {
-    val queue = new mutable.PriorityQueue()(CubeIdOrdering)
+    val queue = new mutable.PriorityQueue()(CubeIdRollupOrdering)
     groups.keys.foreach(queue.enqueue(_))
     while (queue.nonEmpty) {
       val cubeId = queue.dequeue()
@@ -78,11 +78,12 @@ private[writer] class Rollup(limit: Double) {
     }.toMap
   }
 
-  private def areSiblings(cube_a: CubeId, cube_b: CubeId): Boolean = {
-    val sameParent = cube_a.parent == cube_b.parent
-    val differentCube = cube_a != cube_b
-    sameParent && differentCube
-  }
+  /**
+   * Checks if the given cube identifiers are siblings. Two cube identifiers are siblings if they
+   * have the same parent. It is assumed that the cube identifiers are different.
+   */
+  private def areSiblings(cube_a: CubeId, cube_b: CubeId): Boolean =
+    cube_a.parent == cube_b.parent
 
   /*
    * Ordering for cube identifiers. The cube identifiers are ordered by their depth in ascending
@@ -96,7 +97,7 @@ private[writer] class Rollup(limit: Double) {
    * c00, c01, c10, c11, c0, c1, root.
    * c00 -> c01 -> c0, c10 -> c11 -> c1, c0 -> c1 -> root
    */
-  private object CubeIdOrdering extends Ordering[CubeId] {
+  private[writer] object CubeIdRollupOrdering extends Ordering[CubeId] {
 
     override def compare(x: CubeId, y: CubeId): Int = {
       val depthComparison = x.depth.compareTo(y.depth)

diff --git a/src/test/scala/io/qbeast/spark/writer/RollupTest.scala b/src/test/scala/io/qbeast/spark/writer/RollupTest.scala
@@ -19,6 +19,8 @@ import io.qbeast.core.model.CubeId
 import org.scalatest.flatspec.AnyFlatSpec
 import org.scalatest.matchers.should.Matchers
 
+import scala.collection.mutable
+
 /**
  * Tests of Rollup.
  */
@@ -82,4 +84,33 @@ class RollupTest extends AnyFlatSpec with Matchers {
     result(root) shouldBe root
   }
 
+  it should "order CubeIds correctly" in {
+    val root = CubeId.root(4)
+    val Seq(c0, c1, c2) = root.children.take(3).toSeq
+    val Seq(c00, c01, c02, c03) = c0.children.take(4).toSeq
+    val Seq(c10, c11, c12, c13) = c1.children.take(4).toSeq
+    val Seq(c20, c21, c22, c23) = c2.children.take(4).toSeq
+
+    val rollupCubeOrdering = new Rollup(0).CubeIdRollupOrdering
+    val queue = new mutable.PriorityQueue()(rollupCubeOrdering)
+
+    queue.enqueue(root, c0, c1, c2, c00, c01, c02, c03, c10, c11, c12, c13, c20, c21, c22, c23)
+    queue.dequeue() shouldBe c00
+    queue.dequeue() shouldBe c01
+    queue.dequeue() shouldBe c02
+    queue.dequeue() shouldBe c03
+    queue.dequeue() shouldBe c10
+    queue.dequeue() shouldBe c11
+    queue.dequeue() shouldBe c12
+    queue.dequeue() shouldBe c13
+    queue.dequeue() shouldBe c20
+    queue.dequeue() shouldBe c21
+    queue.dequeue() shouldBe c22
+    queue.dequeue() shouldBe c23
+    queue.dequeue() shouldBe c0
+    queue.dequeue() shouldBe c1
+    queue.dequeue() shouldBe c2
+    queue.dequeue() shouldBe root
+  }
+
 }