Spark 开窗源码讲解（二） rownumber

大家好，欢迎来到IT知识分享网。

我们平常使用开窗函数会使用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