Spark 开窗源码讲解(二) rownumber

我们平常使用开窗函数会使用row_number这个函数，目的是根据某个字段分区，然后根据某个字段排序，然后底层的逻辑是怎么样的。

还是看WindowExec 作为入口，然后看windowFrameExpressionFactoryPairs 追溯到它的父类WindowExecBase，看到内部的定义，然后用户可以在代码编辑中使用row_number函数，然后打断点。

可以定位到对应的frame

走的是UnboundedPrecedingWindowFunctionFrame，然后关键核心处就是这个排序。

frameType:rowframe
createBoundOrdering(frameType, upper, timeZone)

然后看这个实现：

说明对应的逻辑在 RowBoundOrdering(0)

private[window] final case class RowBoundOrdering(offset: Int) extends BoundOrdering {
  override def compare(
      inputRow: InternalRow,
      inputIndex: Int,
      outputRow: InternalRow,
      outputIndex: Int): Int =
    inputIndex - (outputIndex + offset)
}

然后看UnboundedPrecedingWindowFunctionFrame的write，在write中就是进行产生这个行号

，可以看到这里的compare方法，这里的 index 就是 上游WindowExecBase中的传入值。

var rowIndex = 0

override final def hasNext: Boolean =
  (bufferIterator != null && bufferIterator.hasNext) || nextRowAvailable

val join = new JoinedRow
override final def next(): InternalRow = {
  // Load the next partition if we need to.
  if ((bufferIterator == null || !bufferIterator.hasNext) && nextRowAvailable) {
    fetchNextPartition()
  }

  if (bufferIterator.hasNext) {
    val current = bufferIterator.next()

    // Get the results for the window frames.
    var i = 0
    while (i < numFrames) {
      frames(i).write(rowIndex, current)
      i += 1
    }

    // 'Merge' the input row with the window function result
    join(current, windowFunctionResult)
    rowIndex += 1

    // Return the projection.
    result(join)
  } else {
    throw new NoSuchElementException
  }
}

这里的 index = outputIndex ,看它的定义，就是逐行遍历增加。

然后看 inputIndex，这个值就是 这个分区内这个组中的当前行的序号，不是分区开始的序号
index是 分区内的序号
inputIndex就是分区内某个组（partitionBy 后面字段的内容的值）的序号。

现在模拟下场景。

+---+---+
| id| kk|
+---+---+
|  1|  1|
|  1|  9|
|  1|  3|
|  1|  1|
+---+---+
这是明细，然后根据 ID进行partitionBy，根据kk排序
第一步：重分区，根据hashpartition 分，ID=1的都分入了同一个分区中
第二步：根据 ID + kk 进行排序，这样数据的顺序从
+---+---+
| id| kk|
+---+---+
|  1|  1|
|  1|  9|
|  1|  3|
|  1|  1|
+---+---+
变成了
+---+---+
| id| kk|
+---+---+
|  1|  1|
|  1|  1|
|  1|  3|
|  1|  9|
+---+---+
第三步：
开始遍历这个分区内部的数据
然后调用 frames(i).write(rowIndex, current)
-------------
override def write(index: Int, current: InternalRow): Unit = {
    var bufferUpdated = index == 0

    // Add all rows to the aggregates for which the input row value is equal to or less than
    // the output row upper bound.
    while (nextRow != null && ubound.compare(nextRow, inputIndex, current, index) <= 0) {
      if (processor != null) {
        processor.update(nextRow)
      }
      nextRow = WindowFunctionFrame.getNextOrNull(inputIterator)
      inputIndex += 1
      bufferUpdated = true
    }

    // Only recalculate and update when the buffer changes.
    if (processor != null && bufferUpdated) {
      processor.evaluate(target)
    }
  }
 此时进入的current 就是（1，1）。nextRow=(1,1),inputIndex=0,index=0,
   ubound.compare(nextRow, inputIndex, current, index) 变成了 inputIndex - (outputIndex + offset)，其中offset=0，等价于 inputIndex - outputIndex，
那么返回的值就是0，0<=0 返回的就是ture.那么processor.update(nextRow)，就产生了一个rownumber=1
此时nextRow 从(1,1) 变成了 (1,1)
然后进入一次循环
此时进入的current就是（1，1），nextRow=(1,1),inputIndex=1,index=1,ubound.compare(nextRow, inputIndex, current, index)返回的依旧是0，
那么processor.update(nextRow)，就产生了一个rownumber=2。
此时nextRow 从(1,1) 变成了 (1,3)
然后进入一次循环
此时进入的current就是（1，3），nextRow=(1,3),产生一个rownumber=3
此时nextRow 从(1,3) 变成了 (1,9)
然后进入一次循环
此时进入的current就是（1，9），nextRow=(1,9),产生一个rownumber=4
此时nextRow 从(1,9) 变成了 null