Tests for NonEmptyRangesCounter

NonEmptyRangesCounterSpec.scala

class NonEmptyRangesCounterSpec extends AnyFlatSpec with Matchers {

  it should "count map-distinct in generated test cases" in {
    import NonEmptyRangesCounterSpec.generateTest
    val totalRanges = 100000
    val totalPartitions = 14
    val maxPositionsPerRange = 20
    val numberOfTests = 10
    Random.setSeed(123456)
    (0 until numberOfTests).foreach { _ =>
      val test = generateTest(totalRanges, maxPositionsPerRange, totalPartitions)
      val result = NonEmptyRangesCounter.count(test.source, test.ranges, test.partitions)
      result should be(test.nonEmptyRanges)
    }
  }
  it should "count non-empty ranges in testcase 1" in {
    val ranges = bitmapOf(0, 15, 16, 19, 23, 29, 47, 80)
    val source = bitmapOf(0, 2, 3, 4, 5, 15, 16, 18, 24, 30, 60, 61, 64, 82, 89, 90)
    val partitions = Seq(100L)

    val distinct = NonEmptyRangesCounter.count(source, ranges, partitions)

    distinct shouldBe Seq(7)
  }

  it should "count non-empty ranges in testcase 2" in {
    val ranges = bitmapOf(0, 5)
    val source = bitmapOf(0, 1, 2, 3, 4)
    val partitions = Seq(10L)

    val distinct = NonEmptyRangesCounter.count(source, ranges, partitions)

    distinct shouldBe Seq(1)
  }

  it should "count non-empty ranges in testcase 3" in {
    val ranges = bitmapOf(0, 5)
    val source = bitmapOf(0, 1, 2, 3, 4, 8)
    val partitions = Seq(10L)

    val distinct = NonEmptyRangesCounter.count(source, ranges, partitions)

    distinct shouldBe Seq(2)
  }

  it should "count non-empty ranges testcase 4" in {
    val ranges = bitmapOf(0, 5, 6, 22, 29, 50)
    val source = bitmapOf(0, 1, 2, 3, 4, 8)
    val partitions = Seq(60L)

    val distinct = NonEmptyRangesCounter.count(source, ranges, partitions)

    distinct shouldBe Seq(2)
  }

  it should "count non-empty ranges when ranges not defined" in {
    val ranges = bitmapOf()
    val source = bitmapOf()
    val partitions = Seq(0L)

    val distinct = NonEmptyRangesCounter.count(source, ranges, partitions)

    distinct shouldBe Seq(0)
  }

  it should "count non-empty ranges when source empty" in {
    val ranges = bitmapOf(0, 4, 6)
    val source = bitmapOf()
    val partitions = Seq(10L)

    val distinct = NonEmptyRangesCounter.count(source, ranges, partitions)

    distinct shouldBe Seq(0)
  }

  it should "count non-empty ranges in partitions" in {
    val ranges = bitmapOf(0, 15, 16, 19, 23, 29, 47, 80)
    val source = bitmapOf(0, 2, 3, 4, 5, 15, 16, 18, 24, 30, 60, 61, 64, 82, 89, 90)
    val partitions = Seq(19L, 47L, 94L)

    val distinct = NonEmptyRangesCounter.count(source, ranges, partitions)

    distinct shouldBe Seq(3, 2, 2)
  }

  it should "count when first partitions are empty" in {
    val ranges = bitmapOf(0, 15, 16, 19, 23, 29, 47, 80)
    val source = bitmapOf(64, 82, 89, 90)
    val partitions = Seq(15L, 19L, 29L, 80L, 94L)

    val distinct = NonEmptyRangesCounter.count(source, ranges, partitions)

    distinct shouldBe Seq(0, 0, 0, 1, 1)
  }

  it should "count when long break in partitions" in {
    val ranges = bitmapOf(0, 15, 16, 19, 23, 29, 47, 80)
    val source = bitmapOf(3, 64, 82, 89, 90)
    val partitions = Seq(15L, 19L, 29L, 80L, 94L)

    val distinct = NonEmptyRangesCounter.count(source, ranges, partitions)

    distinct shouldBe Seq(1, 0, 0, 1, 1)
  }

  it should "count partition" in {
    val ranges = bitmapOf(0, 5)
    val source = bitmapOf(4)
    val partitions = Seq(4L, 6L)

    val distinct = NonEmptyRangesCounter.count(source, ranges, partitions)

    distinct shouldBe Seq(1, 0)
  }

  it should "count partition for huge values" in {
    val ranges = bitmapOf(2, 200, -600)
    val source = bitmapOf(4, -500)
    val ends = Seq(100L, 4294967294L)

    val distinct = NonEmptyRangesCounter.count(source, ranges, ends)

    distinct shouldBe Seq(1, 1)
  }

  it should "pass some regression test" in {
    val ranges = bitmapOf(0, 100)
    val source = bitmapOf(100)
    val partitions = Seq(99L, 100L)

    val distinct = NonEmptyRangesCounter.count(source, ranges, partitions)

    distinct shouldBe Seq(0, 1)
  }

}

object NonEmptyRangesCounterSpec {

  case class Testsource(nonEmptyRanges: Seq[Int],
                        ranges: ImmutableRoaringBitmap,
                        source: ImmutableRoaringBitmap,
                        partitions: Seq[Long])

  def generateTest(totalRanges: Int, maxPositionsPerRange: Int, totalPartitions: Int): Testsource = {
    val totalPositionsOfRange: Map[Int, Int] =
      (0 until totalRanges).map(range => range -> (1 + Random.nextInt(maxPositionsPerRange))).toMap
    val ranges = generateRanges(totalRanges, totalPositionsOfRange)
    val partitions = generatePartitions(ranges, totalPartitions, totalRanges, totalPositionsOfRange)
    val (source, nonEmptyRanges) = generateSource(totalRanges, totalPositionsOfRange, partitions)
    Testsource(nonEmptyRanges, immutable(ranges), immutable(source), partitions)
  }

  private def generateRanges(totalRanges: Int, totalPositionsOfRange: Map[Int, Int]): RoaringBitmap = {
    var pos = 0
    val fun = new RoaringBitmap()
    for (c <- 0 until totalRanges) {
      fun.add(pos)
      pos += totalPositionsOfRange(c)
    }
    fun
  }

  private def generatePartitions(ranges: RoaringBitmap,
                                 totalPartitions: Int,
                                 totalRanges: Int,
                                 totalPositionsOfRange: Map[Int, Int]): Seq[Long] = {
    val step = ranges.last() / totalPartitions
    val it = ranges.getIntIterator
    (0 until totalPartitions).map { partitionNo =>
      it.advanceIfNeeded((partitionNo + 1) * step)
      val isLast = partitionNo == totalPartitions - 1
      if (isLast) {
        Integer.toUnsignedLong(ranges.last()) + totalPositionsOfRange(totalRanges - 1) - 1
      } else {
        Integer.toUnsignedLong(it.next()) - 1
      }
    }
  }

  private def generateSource(totalRanges: Int,
                             totalPositionsOfRange: Map[Int, Int],
                             partitions: Seq[Long]): (RoaringBitmap, Seq[Int]) = {
    val source = new RoaringBitmap()
    var nonEmptyRanges = 0
    var pos = 0L
    val result = Array.ofDim[Int](partitions.length)
    var partNo = 0
    var partitonEnd = partitions(partNo)
    for (c <- 0 until totalRanges) {
      val minReq = pos
      val len = totalPositionsOfRange(c)
      val maxReq = minReq + len - 1
      pos = maxReq + 1

      val addThisRange = Random.nextBoolean()
      if (addThisRange) {
        val includeWholeRange = Random.nextInt(10) == 1
        if (includeWholeRange) {
          source.add(minReq.toLong, maxReq.toLong + 1)
          nonEmptyRanges += 1
        } else {
          var added = false
          for (p <- minReq until maxReq) {
            if (Random.nextBoolean()) {
              source.add(p.toInt)
              added = true
            }
          }
          nonEmptyRanges += (if (added) 1 else 0)
        }
      }
      if (maxReq == partitonEnd) {
        result(partNo) = nonEmptyRanges
        nonEmptyRanges = 0
        partNo += 1
        if (partNo != partitions.length) {
          partitonEnd = partitions(partNo)
        }

      }
    }
    source.runOptimize()
    (source, result.toSeq)
  }

  private def immutable(rb: RoaringBitmap): ImmutableRoaringBitmap = rb.toMutableRoaringBitmap.toImmutableRoaringBitmap
}