SparkSQL自定义bitmap udf操作函数

简介

Bitmap是用来实现基数统计的一种常用方法，它可以实现精确的基数统计。为了提高bitmap对稀疏位图的压缩率，提出了RoaringBitmap（RBM）对稀疏位图进行压缩，减少了内存占用并提高了使用效率。

在Java中实现RBM的常用库是RoaringBitmap，其已经在Spark、Kylin和Druid等系统中得到了应用。那是不是可以将RoaringBitmap库封装成SparkSQL的udf函数，从而可以对bitmap结构数据进行方便的计算操作，实现精确快速地基数统计。

源码实现

参考HLLFunctions的源码实现了三个函数功能：a. rbm_init，初始构造一个RoaringBitmap中间数据，其只支持对Int型的数据进行构建；b. rbm_merge，对rbm_init构造的中间数据进行聚合得到一个聚合后的中间数据；c. rbm_cardinality，求RoaringBitmap中间数据的基数值。

import java.io.{
  ByteArrayInputStream,
  ByteArrayOutputStream,
  DataInputStream,
  DataOutputStream
}

import org.apache.spark.sql.EncapsulationViolator.createAnalysisException
import org.apache.spark.sql.catalyst.{FunctionIdentifier, InternalRow}
import org.apache.spark.sql.catalyst.analysis.{
  FunctionRegistry,
  TypeCheckResult
}
import org.apache.spark.sql.catalyst.expressions.aggregate.{
  AggregateFunction,
  TypedImperativeAggregate
}
import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
import org.apache.spark.sql.catalyst.expressions.{
  ExpectsInputTypes,
  Expression,
  ExpressionDescription,
  ExpressionInfo,
  RuntimeReplaceable,
  UnaryExpression
}
import org.apache.spark.sql.functions.col
import org.apache.spark.sql.types._
import org.apache.spark.sql.{Column, SparkSession}
import org.roaringbitmap.RoaringBitmap

import scala.reflect.ClassTag
import scala.util.{Failure, Success, Try}

trait NullableSketchAggregation
    extends TypedImperativeAggregate[Option[RoaringBitmap]] {

  override def createAggregationBuffer(): Option[RoaringBitmap] = None

  override def merge(buffer: Option[RoaringBitmap],
                     other: Option[RoaringBitmap]): Option[RoaringBitmap] =
    (buffer, other) match {
      case (Some(a), Some(b)) =>
        a.or(b)
        Some(a)
      case (a, None) => a
      case (None, b) => b
      case _         => None
    }

  override def eval(buffer: Option[RoaringBitmap]): Any = {