spark-bitset-sample.scala

val ranges = collect_set(struct($"from", $"to")).as("from_to")
df.groupBy($"id")
  .agg(ranges)
  .withColumn("bytes_sum_unique", Util.findUniqueBytesUDF($"from_to"))
 
// Use the BitSet from java.util.BitSet() due to performance
val findUniqueBytesUDF: UserDefinedFunction = udf { ranges: Seq[Row] =>
  ranges
    .map(x => (x.getAs[Int]("legit_from"), x.getAs[Int]("legit_to")))
    .aggregate(new java.util.BitSet())((bitset, range) => {
      bitset.set(range._1, range._2 + 1)
      bitset
    }, (bitset1, bitset2) => {
      bitset1.or(bitset2)
      bitset1
    })
    .cardinality
}