Rebasing the changes form Qbeast-io#455

Author: cugni

core/src/main/scala/io/qbeast/spark/writer/Rollup.scala

@@ -60,12 +60,15 @@ private[writer] class Rollup(limit: Double) {
       val cubeId = queue.dequeue()
       val group = groups(cubeId)
       if (group.size < limit && !cubeId.isRoot) {
-        val Some(parentCubeId) = cubeId.parent
-        if (groups.contains(parentCubeId)) {
-          groups(parentCubeId).add(group)
+        val nextInLine = cubeId.nextSibling match {
+          case Some(a) => a
+          case None => cubeId.parent.get
+        }
+        if (groups.contains(nextInLine)) {
+          groups(nextInLine).add(group)
         } else {
-          groups.put(parentCubeId, group)
-          queue.enqueue(parentCubeId)
+          groups.put(nextInLine, group)
+          queue.enqueue(nextInLine)
         }
         groups.remove(cubeId)
       }

src/test/scala/io/qbeast/spark/delta/DeltaRollupDataWriterTest.scala

@@ -71,7 +71,7 @@ class DeltaRollupDataWriterTest extends QbeastIntegrationTestSpec {
     }
 
   it should "compute rollup correctly when optimizing" in
-    withSparkAndTmpDir { (spark, tmpDir) =>
+    withSparkAndTmpDir { (_, tmpDir) =>
       val revision =
         Revision(1L, 0, QTableID(tmpDir), 20, Vector(EmptyTransformer("col_1")), Vector.empty)
 
@@ -86,7 +86,7 @@ class DeltaRollupDataWriterTest extends QbeastIntegrationTestSpec {
         Map(root -> 20L, c1 -> 1L, c2 -> 20L))
 
       val rollup = DeltaRollupDataWriter.computeRollup(tc)
-      rollup shouldBe Map(root -> root, c1 -> root, c2 -> c2)
+      rollup shouldBe Map(root -> root, c1 -> c2, c2 -> c2)
     }
 
 }

src/test/scala/io/qbeast/spark/writer/RollupTest.scala

@@ -24,7 +24,7 @@ import org.scalatest.matchers.should.Matchers
  */
 class RollupTest extends AnyFlatSpec with Matchers {
 
-  "Rollup" should "work correctly" in {
+  "Rollup" should "work correctly with basic cube structure" in {
     val root = CubeId.root(1)
     val c0 = root.firstChild
     val c1 = c0.nextSibling.get
@@ -44,8 +44,37 @@ class RollupTest extends AnyFlatSpec with Matchers {
     result(root) shouldBe root
     result(c00) shouldBe c0
     result(c01) shouldBe c0
-    result(c10) shouldBe root
+    result(c10) shouldBe c11 // rolliing up into the next siblings.
     result(c11) shouldBe c11
   }
 
+  it should "handle empty rollup" in {
+    val result = new Rollup(3).compute()
+    result shouldBe empty
+  }
+
+  it should "handle single cube" in {
+    val root = CubeId.root(1)
+    val result = new Rollup(3)
+      .populate(root, 2)
+      .compute()
+
+    result(root) shouldBe root
+  }
+
+  it should "roll up to parent when size exceeds limit" in {
+    val root = CubeId.root(1)
+    val kids = root.children.toSeq
+    val child = kids(0)
+    val grandChild = kids(1)
+
+    val result = new Rollup(2)
+      .populate(root, 1)
+      .populate(child, 2)
+      .populate(grandChild, 3) // Exceeds limit
+      .compute()
+
+    result(grandChild) shouldBe grandChild
+  }
+
 }