Github user ueshin commented on a diff in the pull request: https://github.com/apache/spark/pull/19872#discussion_r157948426 --- Diff: sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFs.scala --- @@ -48,29 +48,46 @@ object ExtractPythonUDFFromAggregate extends Rule[LogicalPlan] { }.isDefined } + private def isPandasGroupAggUdf(expr: Expression): Boolean = expr match { + case PythonUDF(_, _, _, _, PythonEvalType.SQL_PANDAS_GROUP_AGG_UDF) => true + case Alias(child, _) => isPandasGroupAggUdf(child) + case _ => false + } + + private def hasPandasGroupAggUdf(agg: Aggregate): Boolean = { + val actualAggExpr = agg.aggregateExpressions.drop(agg.groupingExpressions.length) + actualAggExpr.exists(isPandasGroupAggUdf) + } + + --- End diff -- nit: remove an extra line.
--- --------------------------------------------------------------------- To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org