HyukjinKwon commented on code in PR #38624:
URL: https://github.com/apache/spark/pull/38624#discussion_r1375549715
##########
python/pyspark/sql/pandas/group_ops.py:
##########
@@ -354,6 +356,132 @@ def applyInPandasWithState(
)
return DataFrame(jdf, self.session)
+ def applyInArrow(
+ self, func: "ArrowGroupedMapFunction", schema: Union[StructType, str]
+ ) -> "DataFrame":
+ """
+ Maps each group of the current :class:`DataFrame` using an Arrow udf
and returns the result
+ as a `DataFrame`.
+
+ The function should take a `pyarrow.Table` and return another
+ `pyarrow.Table`. Alternatively, the user can pass a function that takes
+ a tuple of `pyarrow.Scalar` grouping key(s) and a `pyarrow.Table`.
+ For each group, all columns are passed together as a `pyarrow.Table`
+ to the user-function and the returned `pyarrow.Table` are combined as a
+ :class:`DataFrame`.
+
+ The `schema` should be a :class:`StructType` describing the schema of
the returned
+ `pyarrow.Table`. The column labels of the returned `pyarrow.Table`
must either match
+ the field names in the defined schema if specified as strings, or
match the
+ field data types by position if not strings, e.g. integer indices.
+ The length of the returned `pyarrow.Table` can be arbitrary.
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ func : function
+ a Python native function that takes a `pyarrow.Table` and outputs a
+ `pyarrow.Table`, or that takes one tuple (grouping keys) and a
+ `pyarrow.Table` and outputs a `pyarrow.Table`.
+ schema : :class:`pyspark.sql.types.DataType` or str
+ the return type of the `func` in PySpark. The value can be either a
+ :class:`pyspark.sql.types.DataType` object or a DDL-formatted type
string.
+
+ Examples
+ --------
+ >>> from pyspark.sql.functions import ceil
+ >>> import pyarrow # doctest: +SKIP
+ >>> import pyarrow.compute as pc # doctest: +SKIP
+ >>> df = spark.createDataFrame(
+ ... [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)],
+ ... ("id", "v")) # doctest: +SKIP
+ >>> def normalize(table):
+ ... v = table.column("v")
+ ... norm = pc.divide(pc.subtract(v, pc.mean(v)), pc.stddev(v,
ddof=1))
+ ... return table.set_column(1, "v", norm)
+ >>> df.groupby("id").applyInArrow(
+ ... normalize, schema="id long, v double").show() # doctest: +SKIP
+ +---+-------------------+
+ +---+-------------------+
+ | id| v|
+ +---+-------------------+
+ | 1|-0.7071067811865475|
+ | 1| 0.7071067811865475|
+ | 2|-0.8320502943378437|
+ | 2|-0.2773500981126146|
+ | 2| 1.1094003924504583|
+ +---+-------------------+
+
+ Alternatively, the user can pass a function that takes two arguments.
+ In this case, the grouping key(s) will be passed as the first argument
and the data will
+ be passed as the second argument. The grouping key(s) will be passed
as a tuple of Arrow
+ scalars types, e.g., `pyarrow.Int32Scalar` and `pyarrow.FloatScalar`.
The data will still
+ be passed in as a `pyarrow.Table` containing all columns from the
original Spark DataFrame.
+ This is useful when the user does not want to hardcode grouping key(s)
in the function.
Review Comment:
Should add a newline here; otherwise, the doc would be malformed.
--
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]
