SparkSQL自定义bitmap udf操作函数

最新推荐文章于 2024-09-13 08:10:22 发布

涛声依旧（竞涛）

最新推荐文章于 2024-09-13 08:10:22 发布

阅读量2.7k

点赞数 2

分类专栏： spark 文章标签： SparkSQL bitmap roaringbitmap 源码 udf

本文链接：https://blog.csdn.net/qq_27639777/article/details/110368117

版权

本文介绍了如何将RoaringBitmap封装为SparkSQL的用户自定义函数（UDF），以实现精确的基数统计。通过源码实现，包括rbm_init、rbm_merge和rbm_cardinality三个函数，来处理和聚合bitmap数据。并在SparkSQL中展示了使用这些函数的示例。

摘要由CSDN通过智能技术生成

简介

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 = {
    buffer
      .map(rbm => {
        val bos = new ByteArrayOutputStream
        rbm.serialize(new DataOutputStream(bos))
        bos.toByteArray
      })
      .orNull
  }

  def child: Expressio