You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by gu...@apache.org on 2021/09/20 01:37:35 UTC
[spark] branch master updated: [SPARK-36710][PYTHON] Support new
typing syntax in function apply APIs in pandas API on Spark
This is an automated email from the ASF dual-hosted git repository.
gurwls223 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 8d8b4aa [SPARK-36710][PYTHON] Support new typing syntax in function apply APIs in pandas API on Spark
8d8b4aa is described below
commit 8d8b4aabc3c0b05c837b06335e43e0b5efcdc338
Author: Hyukjin Kwon <gu...@apache.org>
AuthorDate: Mon Sep 20 10:37:06 2021 +0900
[SPARK-36710][PYTHON] Support new typing syntax in function apply APIs in pandas API on Spark
### What changes were proposed in this pull request?
This PR proposes the new syntax introduced in https://github.com/apache/spark/pull/33954. Namely, users now can specify the index type and name as below:
```python
import pandas as pd
import pyspark.pandas as ps
def transform(pdf) -> pd.DataFrame[int, [int, int]]:
pdf['A'] = pdf.id + 1
return pdf
ps.range(5).koalas.apply_batch(transform)
```
```
c0 c1
0 0 1
1 1 2
2 2 3
3 3 4
4 4 5
```
```python
import pandas as pd
import pyspark.pandas as ps
def transform(pdf) -> pd.DataFrame[("index", int), [("a", int), ("b", int)]]:
pdf['A'] = pdf.id * pdf.id
return pdf
ps.range(5).koalas.apply_batch(transform)
```
```
a b
index
0 0 0
1 1 1
2 2 4
3 3 9
4 4 16
```
Again, this syntax remains experimental and this is a non-standard way apart from Python standard. We should migrate to proper typing once pandas supports it like `numpy.typing`.
### Why are the changes needed?
The rationale is described in https://github.com/apache/spark/pull/33954. In order to avoid unnecessary computation for default index or schema inference.
### Does this PR introduce _any_ user-facing change?
Yes, this PR affects the following APIs:
- `DataFrame.apply(..., axis=1)`
- `DataFrame.groupby.apply(...)`
- `DataFrame.pandas_on_spark.transform_batch(...)`
- `DataFrame.pandas_on_spark.apply_batch(...)`
Now they can specify the index type with the new syntax below:
```
DataFrame[index_type, [type, ...]]
DataFrame[(index_name, index_type), [(name, type), ...]]
DataFrame[dtype instance, dtypes instance]
DataFrame[(index_name, index_type), zip(names, types)]
```
### How was this patch tested?
Manually tested, and unittests were added.
Closes #34007 from HyukjinKwon/SPARK-36710.
Authored-by: Hyukjin Kwon <gu...@apache.org>
Signed-off-by: Hyukjin Kwon <gu...@apache.org>
---
python/pyspark/pandas/accessors.py | 112 ++++++++++++++++----------
python/pyspark/pandas/frame.py | 56 +++++++------
python/pyspark/pandas/groupby.py | 58 +++++++++----
python/pyspark/pandas/tests/test_dataframe.py | 66 +++++++++++++++
python/pyspark/pandas/tests/test_groupby.py | 29 +++++++
5 files changed, 236 insertions(+), 85 deletions(-)
diff --git a/python/pyspark/pandas/accessors.py b/python/pyspark/pandas/accessors.py
index d71ded1..8d1d4f7 100644
--- a/python/pyspark/pandas/accessors.py
+++ b/python/pyspark/pandas/accessors.py
@@ -18,7 +18,7 @@
pandas-on-Spark specific features.
"""
import inspect
-from typing import Any, Callable, Optional, Tuple, Union, TYPE_CHECKING, cast
+from typing import Any, Callable, Optional, Tuple, Union, TYPE_CHECKING, cast, List
from types import FunctionType
import numpy as np # noqa: F401
@@ -34,6 +34,7 @@ from pyspark.pandas.internal import (
InternalFrame,
SPARK_INDEX_NAME_FORMAT,
SPARK_DEFAULT_SERIES_NAME,
+ SPARK_DEFAULT_INDEX_NAME,
)
from pyspark.pandas.typedef import infer_return_type, DataFrameType, ScalarType, SeriesType
from pyspark.pandas.utils import (
@@ -211,7 +212,7 @@ class PandasOnSparkFrameMethods(object):
>>> # This case does not return the length of whole frame but of the batch internally
... # used.
- ... def length(pdf) -> ps.DataFrame[int]:
+ ... def length(pdf) -> ps.DataFrame[int, [int]]:
... return pd.DataFrame([len(pdf)])
...
>>> df = ps.DataFrame({'A': range(1000)})
@@ -230,28 +231,24 @@ class PandasOnSparkFrameMethods(object):
To avoid this, specify return type in ``func``, for instance, as below:
- >>> def plus_one(x) -> ps.DataFrame[float, float]:
+ >>> def plus_one(x) -> ps.DataFrame[int, [float, float]]:
... return x + 1
If the return type is specified, the output column names become
`c0, c1, c2 ... cn`. These names are positionally mapped to the returned
DataFrame in ``func``.
- To specify the column names, you can assign them in a pandas friendly style as below:
+ To specify the column names, you can assign them in a NumPy compound type style
+ as below:
- >>> def plus_one(x) -> ps.DataFrame["a": float, "b": float]:
+ >>> def plus_one(x) -> ps.DataFrame[("index", int), [("a", float), ("b", float)]]:
... return x + 1
>>> pdf = pd.DataFrame({'a': [1, 2, 3], 'b': [3, 4, 5]})
- >>> def plus_one(x) -> ps.DataFrame[zip(pdf.dtypes, pdf.columns)]:
+ >>> def plus_one(x) -> ps.DataFrame[
+ ... (pdf.index.name, pdf.index.dtype), zip(pdf.dtypes, pdf.columns)]:
... return x + 1
- When the given function has the return type annotated, the original index of the
- DataFrame will be lost and a default index will be attached to the result DataFrame.
- Please be careful about configuring the default index. See also `Default Index Type
- <https://koalas.readthedocs.io/en/latest/user_guide/options.html#default-index-type>`_.
-
-
Parameters
----------
func : function
@@ -284,17 +281,18 @@ class PandasOnSparkFrameMethods(object):
1 3 4
2 5 6
- >>> def query_func(pdf) -> ps.DataFrame[int, int]:
+ >>> def query_func(pdf) -> ps.DataFrame[int, [int, int]]:
... return pdf.query('A == 1')
>>> df.pandas_on_spark.apply_batch(query_func)
c0 c1
0 1 2
- >>> def query_func(pdf) -> ps.DataFrame["A": int, "B": int]:
+ >>> def query_func(pdf) -> ps.DataFrame[("idx", int), [("A", int), ("B", int)]]:
... return pdf.query('A == 1')
- >>> df.pandas_on_spark.apply_batch(query_func)
- A B
- 0 1 2
+ >>> df.pandas_on_spark.apply_batch(query_func) # doctest: +NORMALIZE_WHITESPACE
+ A B
+ idx
+ 0 1 2
You can also omit the type hints so pandas-on-Spark infers the return schema as below:
@@ -304,7 +302,7 @@ class PandasOnSparkFrameMethods(object):
You can also specify extra arguments.
- >>> def calculation(pdf, y, z) -> ps.DataFrame[int, int]:
+ >>> def calculation(pdf, y, z) -> ps.DataFrame[int, [int, int]]:
... return pdf ** y + z
>>> df.pandas_on_spark.apply_batch(calculation, args=(10,), z=20)
c0 c1
@@ -386,22 +384,32 @@ class PandasOnSparkFrameMethods(object):
"The given function should specify a frame as its type "
"hints; however, the return type was %s." % return_sig
)
+ index_field = cast(DataFrameType, return_type).index_field
+ should_retain_index = index_field is not None
return_schema = cast(DataFrameType, return_type).spark_type
output_func = GroupBy._make_pandas_df_builder_func(
- self_applied, func, return_schema, retain_index=False
+ self_applied, func, return_schema, retain_index=should_retain_index
)
sdf = self_applied._internal.to_internal_spark_frame.mapInPandas(
lambda iterator: map(output_func, iterator), schema=return_schema
)
- # Otherwise, it loses index.
+ index_spark_columns = None
+ index_names: Optional[List[Optional[Tuple[Any, ...]]]] = None
+ index_fields = None
+ if should_retain_index:
+ index_spark_columns = [scol_for(sdf, index_field.struct_field.name)]
+ index_fields = [index_field]
+ if index_field.struct_field.name != SPARK_DEFAULT_INDEX_NAME:
+ index_names = [(index_field.struct_field.name,)]
internal = InternalFrame(
spark_frame=sdf,
- index_spark_columns=None,
- data_fields=cast(DataFrameType, return_type).fields,
+ index_names=index_names,
+ index_spark_columns=index_spark_columns,
+ index_fields=index_fields,
+ data_fields=cast(DataFrameType, return_type).data_fields,
)
-
return DataFrame(internal)
def transform_batch(
@@ -439,28 +447,24 @@ class PandasOnSparkFrameMethods(object):
To avoid this, specify return type in ``func``, for instance, as below:
- >>> def plus_one(x) -> ps.DataFrame[float, float]:
+ >>> def plus_one(x) -> ps.DataFrame[int, [float, float]]:
... return x + 1
If the return type is specified, the output column names become
`c0, c1, c2 ... cn`. These names are positionally mapped to the returned
DataFrame in ``func``.
- To specify the column names, you can assign them in a pandas friendly style as below:
+ To specify the column names, you can assign them in a NumPy compound type style
+ as below:
- >>> def plus_one(x) -> ps.DataFrame['a': float, 'b': float]:
+ >>> def plus_one(x) -> ps.DataFrame[("index", int), [("a", float), ("b", float)]]:
... return x + 1
>>> pdf = pd.DataFrame({'a': [1, 2, 3], 'b': [3, 4, 5]})
- >>> def plus_one(x) -> ps.DataFrame[zip(pdf.dtypes, pdf.columns)]:
+ >>> def plus_one(x) -> ps.DataFrame[
+ ... (pdf.index.name, pdf.index.dtype), zip(pdf.dtypes, pdf.columns)]:
... return x + 1
- When the given function returns DataFrame and has the return type annotated, the
- original index of the DataFrame will be lost and then a default index will be attached
- to the result. Please be careful about configuring the default index. See also
- `Default Index Type
- <https://koalas.readthedocs.io/en/latest/user_guide/options.html#default-index-type>`_.
-
Parameters
----------
func : function
@@ -488,7 +492,7 @@ class PandasOnSparkFrameMethods(object):
1 3 4
2 5 6
- >>> def plus_one_func(pdf) -> ps.DataFrame[int, int]:
+ >>> def plus_one_func(pdf) -> ps.DataFrame[int, [int, int]]:
... return pdf + 1
>>> df.pandas_on_spark.transform_batch(plus_one_func)
c0 c1
@@ -496,13 +500,14 @@ class PandasOnSparkFrameMethods(object):
1 4 5
2 6 7
- >>> def plus_one_func(pdf) -> ps.DataFrame['A': int, 'B': int]:
+ >>> def plus_one_func(pdf) -> ps.DataFrame[("index", int), [('A', int), ('B', int)]]:
... return pdf + 1
- >>> df.pandas_on_spark.transform_batch(plus_one_func)
- A B
- 0 2 3
- 1 4 5
- 2 6 7
+ >>> df.pandas_on_spark.transform_batch(plus_one_func) # doctest: +NORMALIZE_WHITESPACE
+ A B
+ index
+ 0 2 3
+ 1 4 5
+ 2 6 7
>>> def plus_one_func(pdf) -> ps.Series[int]:
... return pdf.B + 1
@@ -551,6 +556,7 @@ class PandasOnSparkFrameMethods(object):
spec = inspect.getfullargspec(func)
return_sig = spec.annotations.get("return", None)
should_infer_schema = return_sig is None
+ should_retain_index = should_infer_schema
original_func = func
func = lambda o: original_func(o, *args, **kwargs)
@@ -672,16 +678,22 @@ class PandasOnSparkFrameMethods(object):
)
return first_series(DataFrame(internal))
else:
+ index_field = cast(DataFrameType, return_type).index_field
+ index_field = (
+ index_field.normalize_spark_type() if index_field is not None else None
+ )
data_fields = [
field.normalize_spark_type()
- for field in cast(DataFrameType, return_type).fields
+ for field in cast(DataFrameType, return_type).data_fields
]
- return_schema = StructType([field.struct_field for field in data_fields])
+ normalized_fields = ([index_field] if index_field is not None else []) + data_fields
+ return_schema = StructType([field.struct_field for field in normalized_fields])
+ should_retain_index = index_field is not None
self_applied = DataFrame(self._psdf._internal.resolved_copy)
output_func = GroupBy._make_pandas_df_builder_func(
- self_applied, func, return_schema, retain_index=False
+ self_applied, func, return_schema, retain_index=should_retain_index
)
columns = self_applied._internal.spark_columns
@@ -693,8 +705,20 @@ class PandasOnSparkFrameMethods(object):
sdf = self_applied._internal.spark_frame.select(applied)
sdf = sdf.selectExpr("%s.*" % temp_struct_column)
+ index_spark_columns = None
+ index_names: Optional[List[Optional[Tuple[Any, ...]]]] = None
+ index_fields = None
+ if should_retain_index:
+ index_spark_columns = [scol_for(sdf, index_field.struct_field.name)]
+ index_fields = [index_field]
+ if index_field.struct_field.name != SPARK_DEFAULT_INDEX_NAME:
+ index_names = [(index_field.struct_field.name,)]
internal = InternalFrame(
- spark_frame=sdf, index_spark_columns=None, data_fields=data_fields
+ spark_frame=sdf,
+ index_names=index_names,
+ index_spark_columns=index_spark_columns,
+ index_fields=index_fields,
+ data_fields=data_fields,
)
return DataFrame(internal)
diff --git a/python/pyspark/pandas/frame.py b/python/pyspark/pandas/frame.py
index 9f93d8d..8d37d31 100644
--- a/python/pyspark/pandas/frame.py
+++ b/python/pyspark/pandas/frame.py
@@ -2317,7 +2317,7 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
In case when axis is 1, it requires to specify `DataFrame` or scalar value
with type hints as below:
- >>> def plus_one(x) -> ps.DataFrame[float, float]:
+ >>> def plus_one(x) -> ps.DataFrame[int, [float, float]]:
... return x + 1
If the return type is specified as `DataFrame`, the output column names become
@@ -2326,22 +2326,14 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
To specify the column names, you can assign them in a pandas friendly style as below:
- >>> def plus_one(x) -> ps.DataFrame["a": float, "b": float]:
+ >>> def plus_one(x) -> ps.DataFrame[("index", int), [("a", float), ("b", float)]]:
... return x + 1
>>> pdf = pd.DataFrame({'a': [1, 2, 3], 'b': [3, 4, 5]})
- >>> def plus_one(x) -> ps.DataFrame[zip(pdf.dtypes, pdf.columns)]:
+ >>> def plus_one(x) -> ps.DataFrame[
+ ... (pdf.index.name, pdf.index.dtype), zip(pdf.dtypes, pdf.columns)]:
... return x + 1
- However, this way switches the index type to default index type in the output
- because the type hint cannot express the index type at this moment. Use
- `reset_index()` to keep index as a workaround.
-
- When the given function has the return type annotated, the original index of the
- DataFrame will be lost and then a default index will be attached to the result.
- Please be careful about configuring the default index. See also `Default Index Type
- <https://koalas.readthedocs.io/en/latest/user_guide/options.html#default-index-type>`_.
-
Parameters
----------
func : function
@@ -2436,18 +2428,19 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
In order to specify the types when `axis` is '1', it should use DataFrame[...]
annotation. In this case, the column names are automatically generated.
- >>> def identify(x) -> ps.DataFrame['A': np.int64, 'B': np.int64]:
+ >>> def identify(x) -> ps.DataFrame[('index', int), [('A', np.int64), ('B', np.int64)]]:
... return x
...
- >>> df.apply(identify, axis=1)
- A B
- 0 4 9
- 1 4 9
- 2 4 9
+ >>> df.apply(identify, axis=1) # doctest: +NORMALIZE_WHITESPACE
+ A B
+ index
+ 0 4 9
+ 1 4 9
+ 2 4 9
You can also specify extra arguments.
- >>> def plus_two(a, b, c) -> ps.DataFrame[np.int64, np.int64]:
+ >>> def plus_two(a, b, c) -> ps.DataFrame[np.int64, [np.int64, np.int64]]:
... return a + b + c
...
>>> df.apply(plus_two, axis=1, args=(1,), c=3)
@@ -2469,6 +2462,7 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
spec = inspect.getfullargspec(func)
return_sig = spec.annotations.get("return", None)
should_infer_schema = return_sig is None
+ should_retain_index = should_infer_schema
def apply_func(pdf: pd.DataFrame) -> pd.DataFrame:
pdf_or_pser = pdf.apply(func, axis=axis, args=args, **kwds)
@@ -2501,7 +2495,7 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
return_schema = StructType([field.struct_field for field in index_fields + data_fields])
output_func = GroupBy._make_pandas_df_builder_func(
- self_applied, apply_func, return_schema, retain_index=True
+ self_applied, apply_func, return_schema, retain_index=should_retain_index
)
sdf = self_applied._internal.to_internal_spark_frame.mapInPandas(
lambda iterator: map(output_func, iterator), schema=return_schema
@@ -2515,6 +2509,7 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
return_type = infer_return_type(func)
require_index_axis = isinstance(return_type, SeriesType)
require_column_axis = isinstance(return_type, DataFrameType)
+ index_field = None
if require_index_axis:
if axis != 0:
@@ -2539,7 +2534,9 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
"hints when axis is 1 or 'column'; however, the return type "
"was %s" % return_sig
)
- data_fields = cast(DataFrameType, return_type).fields
+ index_field = cast(DataFrameType, return_type).index_field
+ should_retain_index = index_field is not None
+ data_fields = cast(DataFrameType, return_type).data_fields
return_schema = cast(DataFrameType, return_type).spark_type
else:
# any axis is fine.
@@ -2558,18 +2555,27 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
column_labels = [None]
output_func = GroupBy._make_pandas_df_builder_func(
- self_applied, apply_func, return_schema, retain_index=False
+ self_applied, apply_func, return_schema, retain_index=should_retain_index
)
sdf = self_applied._internal.to_internal_spark_frame.mapInPandas(
lambda iterator: map(output_func, iterator), schema=return_schema
)
- # Otherwise, it loses index.
+ index_spark_columns = None
+ index_names: Optional[List[Optional[Tuple[Any, ...]]]] = None
+ index_fields = None
+ if should_retain_index:
+ index_spark_columns = [scol_for(sdf, index_field.struct_field.name)]
+ index_fields = [index_field]
+ if index_field.struct_field.name != SPARK_DEFAULT_INDEX_NAME:
+ index_names = [(index_field.struct_field.name,)]
internal = InternalFrame(
spark_frame=sdf,
- index_spark_columns=None,
- column_labels=column_labels,
+ index_names=index_names,
+ index_spark_columns=index_spark_columns,
+ index_fields=index_fields,
data_fields=data_fields,
+ column_labels=column_labels,
)
result = DataFrame(internal) # type: "DataFrame"
diff --git a/python/pyspark/pandas/groupby.py b/python/pyspark/pandas/groupby.py
index 3d4db25..213ecac 100644
--- a/python/pyspark/pandas/groupby.py
+++ b/python/pyspark/pandas/groupby.py
@@ -75,6 +75,7 @@ from pyspark.pandas.internal import (
NATURAL_ORDER_COLUMN_NAME,
SPARK_INDEX_NAME_FORMAT,
SPARK_DEFAULT_SERIES_NAME,
+ SPARK_DEFAULT_INDEX_NAME,
)
from pyspark.pandas.missing.groupby import (
MissingPandasLikeDataFrameGroupBy,
@@ -1026,27 +1027,24 @@ class GroupBy(Generic[FrameLike], metaclass=ABCMeta):
To avoid this, specify return type in ``func``, for instance, as below:
- >>> def pandas_div(x) -> ps.DataFrame[float, float]:
+ >>> def pandas_div(x) -> ps.DataFrame[int, [float, float]]:
... return x[['B', 'C']] / x[['B', 'C']]
If the return type is specified, the output column names become
`c0, c1, c2 ... cn`. These names are positionally mapped to the returned
DataFrame in ``func``.
- To specify the column names, you can assign them in a pandas friendly style as below:
+ To specify the column names, you can assign them in a NumPy compound type style
+ as below:
- >>> def pandas_div(x) -> ps.DataFrame["a": float, "b": float]:
+ >>> def pandas_div(x) -> ps.DataFrame[("index", int), [("a", float), ("b", float)]]:
... return x[['B', 'C']] / x[['B', 'C']]
>>> pdf = pd.DataFrame({'B': [1.], 'C': [3.]})
- >>> def plus_one(x) -> ps.DataFrame[zip(pdf.columns, pdf.dtypes)]:
+ >>> def plus_one(x) -> ps.DataFrame[
+ ... (pdf.index.name, pdf.index.dtype), zip(pdf.columns, pdf.dtypes)]:
... return x[['B', 'C']] / x[['B', 'C']]
- When the given function has the return type annotated, the original index of the
- GroupBy object will be lost and a default index will be attached to the result.
- Please be careful about configuring the default index. See also `Default Index Type
- <https://koalas.readthedocs.io/en/latest/user_guide/options.html#default-index-type>`_.
-
.. note:: the dataframe within ``func`` is actually a pandas dataframe. Therefore,
any pandas API within this function is allowed.
@@ -1106,7 +1104,7 @@ class GroupBy(Generic[FrameLike], metaclass=ABCMeta):
You can specify the type hint and prevent schema inference for better performance.
- >>> def pandas_div(x) -> ps.DataFrame[float, float]:
+ >>> def pandas_div(x) -> ps.DataFrame[int, [float, float]]:
... return x[['B', 'C']] / x[['B', 'C']]
>>> g.apply(pandas_div).sort_index() # doctest: +NORMALIZE_WHITESPACE
c0 c1
@@ -1114,6 +1112,15 @@ class GroupBy(Generic[FrameLike], metaclass=ABCMeta):
1 1.0 1.0
2 1.0 1.0
+ >>> def pandas_div(x) -> ps.DataFrame[("index", int), [("f1", float), ("f2", float)]]:
+ ... return x[['B', 'C']] / x[['B', 'C']]
+ >>> g.apply(pandas_div).sort_index() # doctest: +NORMALIZE_WHITESPACE
+ f1 f2
+ index
+ 0 1.0 1.0
+ 1 1.0 1.0
+ 2 1.0 1.0
+
In case of Series, it works as below.
>>> def plus_max(x) -> ps.Series[np.int]:
@@ -1156,6 +1163,7 @@ class GroupBy(Generic[FrameLike], metaclass=ABCMeta):
spec = inspect.getfullargspec(func)
return_sig = spec.annotations.get("return", None)
should_infer_schema = return_sig is None
+ should_retain_index = should_infer_schema
is_series_groupby = isinstance(self, SeriesGroupBy)
@@ -1226,8 +1234,12 @@ class GroupBy(Generic[FrameLike], metaclass=ABCMeta):
)
if isinstance(return_type, DataFrameType):
- data_fields = cast(DataFrameType, return_type).fields
+ data_fields = cast(DataFrameType, return_type).data_fields
return_schema = cast(DataFrameType, return_type).spark_type
+ index_field = cast(DataFrameType, return_type).index_field
+ should_retain_index = index_field is not None
+ index_fields = [index_field]
+ psdf_from_pandas = None
else:
should_return_series = True
dtype = cast(Union[SeriesType, ScalarType], return_type).dtype
@@ -1288,14 +1300,28 @@ class GroupBy(Generic[FrameLike], metaclass=ABCMeta):
pandas_groupby_apply,
[psdf._internal.spark_column_for(label) for label in groupkey_labels],
return_schema,
- retain_index=should_infer_schema,
+ retain_index=should_retain_index,
)
- if should_infer_schema:
+ if should_retain_index:
# If schema is inferred, we can restore indexes too.
- internal = psdf_from_pandas._internal.with_new_sdf(
- spark_frame=sdf, index_fields=index_fields, data_fields=data_fields
- )
+ if psdf_from_pandas is not None:
+ internal = psdf_from_pandas._internal.with_new_sdf(
+ spark_frame=sdf, index_fields=index_fields, data_fields=data_fields
+ )
+ else:
+ index_names: Optional[List[Optional[Tuple[Any, ...]]]] = None
+ index_field = index_fields[0]
+ index_spark_columns = [scol_for(sdf, index_field.struct_field.name)]
+ if index_field.struct_field.name != SPARK_DEFAULT_INDEX_NAME:
+ index_names = [(index_field.struct_field.name,)]
+ internal = InternalFrame(
+ spark_frame=sdf,
+ index_names=index_names,
+ index_spark_columns=index_spark_columns,
+ index_fields=index_fields,
+ data_fields=data_fields,
+ )
else:
# Otherwise, it loses index.
internal = InternalFrame(
diff --git a/python/pyspark/pandas/tests/test_dataframe.py b/python/pyspark/pandas/tests/test_dataframe.py
index d964f9f..67df870 100644
--- a/python/pyspark/pandas/tests/test_dataframe.py
+++ b/python/pyspark/pandas/tests/test_dataframe.py
@@ -4428,6 +4428,28 @@ class DataFrameTest(PandasOnSparkTestCase, SQLTestUtils):
pdf.apply(lambda x: len(x), axis=1).sort_index(),
)
+ def test_apply_with_type(self):
+ pdf = self.pdf
+ psdf = ps.from_pandas(pdf)
+
+ def identify1(x) -> ps.DataFrame[int, int]:
+ return x
+
+ # Type hints set the default column names, and we use default index for
+ # pandas API on Spark. Here we ignore both diff.
+ actual = psdf.apply(identify1, axis=1)
+ expected = pdf.apply(identify1, axis=1)
+ self.assert_eq(sorted(actual["c0"].to_numpy()), sorted(expected["a"].to_numpy()))
+ self.assert_eq(sorted(actual["c1"].to_numpy()), sorted(expected["b"].to_numpy()))
+
+ def identify2(x) -> ps.DataFrame[slice("a", int), slice("b", int)]: # noqa: F405
+ return x
+
+ actual = psdf.apply(identify2, axis=1)
+ expected = pdf.apply(identify2, axis=1)
+ self.assert_eq(sorted(actual["a"].to_numpy()), sorted(expected["a"].to_numpy()))
+ self.assert_eq(sorted(actual["b"].to_numpy()), sorted(expected["b"].to_numpy()))
+
def test_apply_batch(self):
pdf = pd.DataFrame(
{
@@ -4481,6 +4503,28 @@ class DataFrameTest(PandasOnSparkTestCase, SQLTestUtils):
(pdf + 1).sort_index(),
)
+ def test_apply_batch_with_type(self):
+ pdf = self.pdf
+ psdf = ps.from_pandas(pdf)
+
+ def identify1(x) -> ps.DataFrame[int, int]:
+ return x
+
+ # Type hints set the default column names, and we use default index for
+ # pandas API on Spark. Here we ignore both diff.
+ actual = psdf.pandas_on_spark.apply_batch(identify1)
+ expected = pdf
+ self.assert_eq(sorted(actual["c0"].to_numpy()), sorted(expected["a"].to_numpy()))
+ self.assert_eq(sorted(actual["c1"].to_numpy()), sorted(expected["b"].to_numpy()))
+
+ def identify2(x) -> ps.DataFrame[slice("a", int), slice("b", int)]: # noqa: F405
+ return x
+
+ actual = psdf.pandas_on_spark.apply_batch(identify2)
+ expected = pdf
+ self.assert_eq(sorted(actual["a"].to_numpy()), sorted(expected["a"].to_numpy()))
+ self.assert_eq(sorted(actual["b"].to_numpy()), sorted(expected["b"].to_numpy()))
+
def test_transform_batch(self):
pdf = pd.DataFrame(
{
@@ -4550,6 +4594,28 @@ class DataFrameTest(PandasOnSparkTestCase, SQLTestUtils):
(pdf + 1).sort_index(),
)
+ def test_transform_batch_with_type(self):
+ pdf = self.pdf
+ psdf = ps.from_pandas(pdf)
+
+ def identify1(x) -> ps.DataFrame[int, int]:
+ return x
+
+ # Type hints set the default column names, and we use default index for
+ # pandas API on Spark. Here we ignore both diff.
+ actual = psdf.pandas_on_spark.transform_batch(identify1)
+ expected = pdf
+ self.assert_eq(sorted(actual["c0"].to_numpy()), sorted(expected["a"].to_numpy()))
+ self.assert_eq(sorted(actual["c1"].to_numpy()), sorted(expected["b"].to_numpy()))
+
+ def identify2(x) -> ps.DataFrame[slice("a", int), slice("b", int)]: # noqa: F405
+ return x
+
+ actual = psdf.pandas_on_spark.transform_batch(identify2)
+ expected = pdf
+ self.assert_eq(sorted(actual["a"].to_numpy()), sorted(expected["a"].to_numpy()))
+ self.assert_eq(sorted(actual["b"].to_numpy()), sorted(expected["b"].to_numpy()))
+
def test_transform_batch_same_anchor(self):
psdf = ps.range(10)
psdf["d"] = psdf.pandas_on_spark.transform_batch(lambda pdf: pdf.id + 1)
diff --git a/python/pyspark/pandas/tests/test_groupby.py b/python/pyspark/pandas/tests/test_groupby.py
index 0640001..d85c3ab 100644
--- a/python/pyspark/pandas/tests/test_groupby.py
+++ b/python/pyspark/pandas/tests/test_groupby.py
@@ -1987,6 +1987,35 @@ class GroupByTest(PandasOnSparkTestCase, TestUtils):
with option_context("compute.shortcut_limit", 0):
self.test_apply()
+ def test_apply_with_type_hint(self):
+ pdf = pd.DataFrame(
+ {"a": [1, 2, 3, 4, 5, 6], "b": [1, 1, 2, 3, 5, 8], "c": [1, 4, 9, 16, 25, 36]},
+ columns=["a", "b", "c"],
+ )
+ psdf = ps.from_pandas(pdf)
+
+ def add_max1(x) -> ps.DataFrame[int, int, int]:
+ return x + x.min()
+
+ # Type hints set the default column names, and we use default index for
+ # pandas API on Spark. Here we ignore both diff.
+ actual = psdf.groupby("b").apply(add_max1).sort_index()
+ expected = pdf.groupby("b").apply(add_max1).sort_index()
+ self.assert_eq(sorted(actual["c0"].to_numpy()), sorted(expected["a"].to_numpy()))
+ self.assert_eq(sorted(actual["c1"].to_numpy()), sorted(expected["b"].to_numpy()))
+ self.assert_eq(sorted(actual["c2"].to_numpy()), sorted(expected["c"].to_numpy()))
+
+ def add_max2(
+ x,
+ ) -> ps.DataFrame[slice("a", int), slice("b", int), slice("c", int)]: # noqa: F405
+ return x + x.min()
+
+ actual = psdf.groupby("b").apply(add_max2).sort_index()
+ expected = pdf.groupby("b").apply(add_max2).sort_index()
+ self.assert_eq(sorted(actual["a"].to_numpy()), sorted(expected["a"].to_numpy()))
+ self.assert_eq(sorted(actual["c"].to_numpy()), sorted(expected["c"].to_numpy()))
+ self.assert_eq(sorted(actual["c"].to_numpy()), sorted(expected["c"].to_numpy()))
+
def test_apply_negative(self):
def func(_) -> ps.Series[int]:
return pd.Series([1])
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org