icexelloss commented on a change in pull request #22305:
[SPARK-24561][SQL][Python] User-defined window aggregation functions with
Pandas UDF (bounded window)
URL: https://github.com/apache/spark/pull/22305#discussion_r242302249
##########
File path:
sql/core/src/main/scala/org/apache/spark/sql/execution/python/WindowInPandasExec.scala
##########
@@ -73,68 +141,141 @@ case class WindowInPandasExec(
}
/**
- * Create the resulting projection.
- *
- * This method uses Code Generation. It can only be used on the executor
side.
- *
- * @param expressions unbound ordered function expressions.
- * @return the final resulting projection.
+ * See [[WindowBoundHelpers]] for details.
*/
- private[this] def createResultProjection(expressions: Seq[Expression]):
UnsafeProjection = {
- val references = expressions.zipWithIndex.map { case (e, i) =>
- // Results of window expressions will be on the right side of child's
output
- BoundReference(child.output.size + i, e.dataType, e.nullable)
+ private def computeWindowBoundHelpers(
+ factories: Seq[InternalRow => WindowFunctionFrame]
+ ): WindowBoundHelpers = {
+ val dummyRow = EmptyRow
+ val functionFrames = factories.map(_(dummyRow))
+
+ val windowBoundTypes = functionFrames.map {
+ case _: UnboundedWindowFunctionFrame => UnboundedWindow
+ case _: UnboundedFollowingWindowFunctionFrame |
+ _: SlidingWindowFunctionFrame |
+ _: UnboundedPrecedingWindowFunctionFrame => BoundedWindow
+ // It should be impossible to get other types of window function frame
here
+ case frame => throw new RuntimeException(s"Unexpected window function
frame $frame.")
+ }
+
+ val requiredIndices = functionFrames.map {
+ case _: UnboundedWindowFunctionFrame => 0
+ case _ => 2
+ }
+
+ val upperBoundIndices = requiredIndices.scan(0)(_ + _).tail
+
+ val boundIndices = requiredIndices.zip(upperBoundIndices).map { case (num,
upperBoundIndex) =>
+ if (num == 0) {
+ // Sentinel values for unbounded window
+ (-1, -1)
+ } else {
+ (upperBoundIndex - 2, upperBoundIndex - 1)
+ }
}
- val unboundToRefMap = expressions.zip(references).toMap
- val patchedWindowExpression =
windowExpression.map(_.transform(unboundToRefMap))
- UnsafeProjection.create(
- child.output ++ patchedWindowExpression,
- child.output)
+
+ def lowerBoundIndex(frameIndex: Int) = boundIndices(frameIndex)._1
+ def upperBoundIndex(frameIndex: Int) = boundIndices(frameIndex)._2
+
+ (requiredIndices.sum, lowerBoundIndex, upperBoundIndex, windowBoundTypes)
}
protected override def doExecute(): RDD[InternalRow] = {
- val inputRDD = child.execute()
+ // Unwrap the expressions and factories from the map.
+ val expressionsWithFrameIndex =
+ windowFrameExpressionFactoryPairs.map(_._1).zipWithIndex.flatMap {
+ case (buffer, frameIndex) => buffer.map( expr => (expr, frameIndex))
Review comment:
Thanks! Fixed
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
