[Delta][SQL] Delta开源付费功能，最全分析ZOrder的源码实现流程

x value：01100001  97     98
y value：11010110  104    105

z-value： 1011011000101001 46633

OPTIMIZE delta.table WHERE day = 25 AND city = ‘New York’ ZORDER BY (col1, col2)

| OPTIMIZE (path=STRING | table=qualifiedName)
    (WHERE partitionPredicate = predicateToken)?
    (zorderSpec)?                                                   #optimizeTable

zorderSpec
    : ZORDER BY LEFT_PAREN interleave+=qualifiedName (COMMA interleave+=qualifiedName)* RIGHT_PAREN
    | ZORDER BY interleave+=qualifiedName (COMMA interleave+=qualifiedName)*
    ;

override def visitOptimizeTable(ctx: OptimizeTableContext): AnyRef = withOrigin(ctx) {
  if (ctx.path == null && ctx.table == null) {
    throw new ParseException("OPTIMIZE command requires a file path or table name.", ctx)
  }
  // z-order的列的seq
  val interleaveBy = Option(ctx.zorderSpec).map(visitZorderSpec).getOrElse(Seq.empty)
  OptimizeTableCommand(
    Option(ctx.path).map(string),
    Option(ctx.table).map(visitTableIdentifier),
    Option(ctx.partitionPredicate).map(extractRawText(_)))(interleaveBy)
}

override def run(sparkSession: SparkSession): Seq[Row] = {
  val deltaLog = getDeltaLog(sparkSession, path, tableId, "OPTIMIZE")

  // [1] 从metadata中获取表的分区列
  val partitionColumns = deltaLog.snapshot.metadata.partitionColumns
  // [2] 解析查询谓词
  val partitionPredicates = partitionPredicate.map(predicate => {
    val predicates = parsePredicates(sparkSession, predicate)
    verifyPartitionPredicates(
      sparkSession,
      partitionColumns,
      predicates)
    predicates
  }).getOrElse(Seq.empty)
  // [3] 校验zorder列
  validateZorderByColumns(sparkSession, deltaLog, zOrderBy)
  val zOrderByColumns = zOrderBy.map(_.name).toSeq
  // [4] 调用optimize命令
  new OptimizeExecutor(sparkSession, deltaLog, partitionPredicates, zOrderByColumns)
      .optimize()
}

def optimize(): Seq[Row] = {
  val txn = deltaLog.startTransaction()
  // [1] 按照where条件进行筛选出候选文件
  val candidateFiles = txn.filterFiles(partitionPredicate)
  val partitionSchema = txn.metadata.partitionSchema
  // [2] 注意这里如果isMultiDimClustering多维聚集则不过滤文件的大小直接选择所有的文件
  // select all files in case of multi-dimensional clustering
  val filesToProcess = candidateFiles.filter(_.size < minFileSize || isMultiDimClustering)
  val partitionsToCompact = filesToProcess.groupBy(_.partitionValues).toSeq

  val jobs = groupFilesIntoBins(partitionsToCompact, maxFileSize)

val parallelJobCollection = new ParVector(jobs.toVector)

// Create a task pool to parallelize the submission of optimization jobs to Spark.
val threadPool = ThreadUtils.newForkJoinPool(
  "OptimizeJob",
  sparkSession.sessionState.conf.getConf(DeltaSQLConf.DELTA_OPTIMIZE_MAX_THREADS))

val updates = try {
  val forkJoinPoolTaskSupport = new ForkJoinTaskSupport(threadPool)
  parallelJobCollection.tasksupport = forkJoinPoolTaskSupport
  // 并行的执行文件的合并、压缩与Zorder优化
  parallelJobCollection.flatMap(partitionBinGroup =>
    runOptimizeBinJob(txn, partitionBinGroup._1, partitionBinGroup._2, maxFileSize)).seq
} finally {
  threadPool.shutdownNow()
}

// [1] 读取分组的文件
val input = txn.deltaLog.createDataFrame(txn.snapshot, bin, actionTypeOpt = Some("Optimize"))
val repartitionDF = if (isMultiDimClustering) {
  val totalSize = bin.map(_.size).sum
  // 分区数为分组文件的size大小除以其中最大的文件size
  val approxNumFiles = Math.max(1, totalSize / maxFileSize).toInt
  // [2] 调用 MultiDimClustering.cluster
  MultiDimClustering.cluster(
    input,
    approxNumFiles,
    zOrderByColumns)
} else {
  input.coalesce(numPartitions = 1)
}

override def cluster(
    df: DataFrame,
    colNames: Seq[String],
    approxNumPartitions: Int): DataFrame = {
  val conf = df.sparkSession.sessionState.conf
  // 用于控制rangeId，越大，精度越好，但以牺牲性能为代价
  val numRanges = conf.getConf(DeltaSQLConf.MDC_NUM_RANGE_IDS)
  // 是否追加噪音，为了避免数据倾斜的追加的噪音后缀
  val addNoise = conf.getConf(DeltaSQLConf.MDC_ADD_NOISE)

  val cols = colNames.map(df(_))
  // 执行cluster表达式，生成z-value
  val mdcCol = getClusteringExpression(cols, numRanges)
  val repartitionKeyColName = s"${UUID.randomUUID().toString}-rpKey1"
  ...
}

object ZOrderClustering extends SpaceFillingCurveClustering {
  override protected[skipping] def getClusteringExpression(
      cols: Seq[Column], numRanges: Int): Column = {
    assert(cols.size >= 1, "Cannot do Z-Order clustering by zero columns!")
    // [1] 调用range_partition_id表达式生成rangeIdCols
    val rangeIdCols = cols.map(range_partition_id(_, numRanges))
    // [2] 执行interleave_bits，并转换为String
    interleave_bits(rangeIdCols: _*).cast(StringType)
  }
}

def range_partition_id(col: Column, numPartitions: Int): Column = withExpr {
  RangePartitionId(col.expr, numPartitions)
}

extensions.injectOptimizerRule { session =>
  new RangePartitionIdRewrite(session)
}

case RangePartitionId(expr, n) =>
  val aliasedExpr = Alias(expr, "__RPI_child_col__")()
  val exprAttr = aliasedExpr.toAttribute
  // [1] 对于查询列过滤null的行
  val planForSampling = Filter(IsNotNull(exprAttr), Project(Seq(aliasedExpr), node.child))
  ...
  withCallSite(session.sparkContext, desc) {
    SQLExecution.withNewExecutionId(qeForSampling) {
      withJobGroup(session.sparkContext, jobGroupId, desc) {
        // [2] 创建一个pair(InternalRow, null), 用于存储查询列对应的rangeid
        val rddForSampling = qeForSampling.toRdd.mapPartitionsInternal { iter =>
             val mutablePair = new MutablePair[InternalRow, Null]()
             iter.map(row => mutablePair.update(row.copy(), null))
        }
        // [3] 创建RangePartitioner，传入排序的sortOrder
        val sortOrder = SortOrder(exprAttr, Ascending)
        implicit val ordering = new LazilyGeneratedOrdering(Seq(sortOrder), Seq(exprAttr))
        val partitioner = new RangePartitioner(n, rddForSampling, true, sampleSizeHint)
        // [4] 调用PartitionerExpr，执行写入rangeid
        PartitionerExpr(expr, partitioner)

override def eval(input: InternalRow): Any = {
  val value: Any = child.eval(input)
  row.update(0, value)
  partitioner.getPartition(row)
}

var repartitionedDf = if (addNoise) {
  val randByteColName = s"${UUID.randomUUID().toString}-rpKey2"
  val randByteCol = (rand() * 255 - 128).cast(ByteType)
  df.withColumn(repartitionKeyColName, mdcCol).withColumn(randByteColName, randByteCol)
    .repartitionByRange(approxNumPartitions, col(repartitionKeyColName), col(randByteColName))
    .drop(randByteColName)
} else {
  df.withColumn(repartitionKeyColName, mdcCol)
    .repartitionByRange(approxNumPartitions, col(repartitionKeyColName))
}

repartitionedDf.drop(repartitionKeyColName)

if (isMultiDimClustering) {
  val inputFileStats =
    ZOrderFileStats(removedFiles.size, removedFiles.map(_.size.getOrElse(0L)).sum)
  optimizeStats.zOrderStats = Some(ZOrderStats(
    strategyName = "all", // means process all files in a partition
    inputCubeFiles = ZOrderFileStats(0, 0),
    inputOtherFiles = inputFileStats,
    inputNumCubes = 0,
    mergedFiles = inputFileStats,
    // There will one z-cube for each partition
    numOutputCubes = optimizeStats.numPartitionsOptimized))
}