santosh-d3vpl3x commented on code in PR #40122:
URL: https://github.com/apache/spark/pull/40122#discussion_r1114960540
##########
sql/core/src/main/scala/org/apache/spark/sql/RelationalGroupedDataset.scala:
##########
@@ -582,48 +582,59 @@ class RelationalGroupedDataset protected[sql](
}
/**
- * Applies a vectorized python user-defined function to each cogrouped data.
- * The user-defined function defines a transformation:
- * `pandas.DataFrame`, `pandas.DataFrame` -> `pandas.DataFrame`.
- * For each group in the cogrouped data, all elements in the group are
passed as a
- * `pandas.DataFrame` and the results for all cogroups are combined into a
new [[DataFrame]].
+ * Applies a vectorized python user-defined function to each cogrouped data.
The user-defined
+ * function defines a transformation: `pandas.DataFrame`, `pandas.DataFrame`
->
+ * `pandas.DataFrame`. For each group in the cogrouped data, all elements in
the group are
+ * passed as a `pandas.DataFrame` and the results for all cogroups are
combined into a new
+ * [[DataFrame]].
*
* This function uses Apache Arrow as serialization format between Java
executors and Python
* workers.
*/
private[sql] def flatMapCoGroupsInPandas(
- r: RelationalGroupedDataset,
+ rs: Seq[RelationalGroupedDataset],
expr: PythonUDF): DataFrame = {
require(expr.evalType == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF,
"Must pass a cogrouped map udf")
- require(this.groupingExprs.length == r.groupingExprs.length,
- "Cogroup keys must have same size: " +
- s"${this.groupingExprs.length} != ${r.groupingExprs.length}")
+ val groupingExprLengthEquals = rs.map(_.groupingExprs.length).
+ forall(_ == this.groupingExprs.length)
+
+ require(groupingExprLengthEquals, s"Cogroup keys must have same size.")
Review Comment:
Good suggestion, I will add them.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
