You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ru...@apache.org on 2023/12/20 06:59:10 UTC
(spark) branch master updated: [SPARK-46433][PS][TESTS][FOLLOWUPS] Reorganize `OpsOnDiffFramesGroupByTests`: Factor out remaining slow tests
This is an automated email from the ASF dual-hosted git repository.
ruifengz 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 771eb89ab2a4 [SPARK-46433][PS][TESTS][FOLLOWUPS] Reorganize `OpsOnDiffFramesGroupByTests`: Factor out remaining slow tests
771eb89ab2a4 is described below
commit 771eb89ab2a46269559f5e1c88afe13fc958e2f1
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Wed Dec 20 14:58:47 2023 +0800
[SPARK-46433][PS][TESTS][FOLLOWUPS] Reorganize `OpsOnDiffFramesGroupByTests`: Factor out remaining slow tests
### What changes were proposed in this pull request?
Factor out remaining slow tests in `OpsOnDiffFramesGroupByTests`
### Why are the changes needed?
to make it more suitable for parallelism
### Does this PR introduce _any_ user-facing change?
no, test-only
### How was this patch tested?
ci
### Was this patch authored or co-authored using generative AI tooling?
no
Closes #44418 from zhengruifeng/ps_test_diff_groupby_3.
Authored-by: Ruifeng Zheng <ru...@apache.org>
Signed-off-by: Ruifeng Zheng <ru...@apache.org>
---
dev/sparktestsupport/modules.py | 10 +-
.../test_parity_groupby.py} | 10 +-
.../test_parity_groupby_shift.py} | 10 +-
.../test_parity_groupby_split_apply_combine.py} | 10 +-
.../test_parity_groupby_transform.py} | 10 +-
.../pandas/tests/diff_frames_ops/test_groupby.py | 127 +++++++++
.../tests/diff_frames_ops/test_groupby_shift.py | 85 ++++++
.../test_groupby_split_apply_combine.py | 113 ++++++++
.../diff_frames_ops/test_groupby_transform.py | 82 ++++++
.../tests/test_ops_on_diff_frames_groupby.py | 287 ---------------------
10 files changed, 439 insertions(+), 305 deletions(-)
diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py
index 1c632871ba61..80930c22af8b 100644
--- a/dev/sparktestsupport/modules.py
+++ b/dev/sparktestsupport/modules.py
@@ -859,7 +859,6 @@ pyspark_pandas_slow = Module(
"pyspark.pandas.tests.groupby.test_value_counts",
"pyspark.pandas.tests.test_indexing",
"pyspark.pandas.tests.test_ops_on_diff_frames",
- "pyspark.pandas.tests.test_ops_on_diff_frames_groupby",
"pyspark.pandas.tests.diff_frames_ops.test_align",
"pyspark.pandas.tests.diff_frames_ops.test_basic_slow",
"pyspark.pandas.tests.diff_frames_ops.test_cov",
@@ -871,6 +870,7 @@ pyspark_pandas_slow = Module(
"pyspark.pandas.tests.diff_frames_ops.test_series",
"pyspark.pandas.tests.diff_frames_ops.test_setitem_frame",
"pyspark.pandas.tests.diff_frames_ops.test_setitem_series",
+ "pyspark.pandas.tests.diff_frames_ops.test_groupby",
"pyspark.pandas.tests.diff_frames_ops.test_groupby_aggregate",
"pyspark.pandas.tests.diff_frames_ops.test_groupby_apply",
"pyspark.pandas.tests.diff_frames_ops.test_groupby_cumulative",
@@ -878,6 +878,9 @@ pyspark_pandas_slow = Module(
"pyspark.pandas.tests.diff_frames_ops.test_groupby_diff_len",
"pyspark.pandas.tests.diff_frames_ops.test_groupby_fillna",
"pyspark.pandas.tests.diff_frames_ops.test_groupby_filter",
+ "pyspark.pandas.tests.diff_frames_ops.test_groupby_shift",
+ "pyspark.pandas.tests.diff_frames_ops.test_groupby_split_apply_combine",
+ "pyspark.pandas.tests.diff_frames_ops.test_groupby_transform",
"pyspark.pandas.tests.series.test_all_any",
"pyspark.pandas.tests.series.test_arg_ops",
"pyspark.pandas.tests.series.test_as_of",
@@ -1215,7 +1218,7 @@ pyspark_pandas_connect_part3 = Module(
"pyspark.pandas.tests.connect.indexes.test_parity_datetime_property",
"pyspark.pandas.tests.connect.indexes.test_parity_datetime_round",
"pyspark.pandas.tests.connect.test_parity_ops_on_diff_frames",
- "pyspark.pandas.tests.connect.test_parity_ops_on_diff_frames_groupby",
+ "pyspark.pandas.tests.connect.diff_frames_ops.test_parity_groupby",
"pyspark.pandas.tests.connect.diff_frames_ops.test_parity_groupby_aggregate",
"pyspark.pandas.tests.connect.diff_frames_ops.test_parity_groupby_apply",
"pyspark.pandas.tests.connect.diff_frames_ops.test_parity_groupby_cumulative",
@@ -1223,6 +1226,9 @@ pyspark_pandas_connect_part3 = Module(
"pyspark.pandas.tests.connect.diff_frames_ops.test_parity_groupby_diff_len",
"pyspark.pandas.tests.connect.diff_frames_ops.test_parity_groupby_fillna",
"pyspark.pandas.tests.connect.diff_frames_ops.test_parity_groupby_filter",
+ "pyspark.pandas.tests.connect.diff_frames_ops.test_parity_groupby_shift",
+ "pyspark.pandas.tests.connect.diff_frames_ops.test_parity_groupby_split_apply_combine",
+ "pyspark.pandas.tests.connect.diff_frames_ops.test_parity_groupby_transform",
],
excluded_python_implementations=[
"PyPy" # Skip these tests under PyPy since they require numpy, pandas, and pyarrow and
diff --git a/python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames_groupby.py b/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_groupby.py
similarity index 79%
copy from python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames_groupby.py
copy to python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_groupby.py
index 685ec5c45c5f..87cde801d1f3 100644
--- a/python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames_groupby.py
+++ b/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_groupby.py
@@ -16,19 +16,21 @@
#
import unittest
-from pyspark.pandas.tests.test_ops_on_diff_frames_groupby import OpsOnDiffFramesGroupByTestsMixin
+from pyspark.pandas.tests.diff_frames_ops.test_groupby import GroupByMixin
from pyspark.testing.connectutils import ReusedConnectTestCase
from pyspark.testing.pandasutils import PandasOnSparkTestUtils
-class OpsOnDiffFramesGroupByParityTests(
- OpsOnDiffFramesGroupByTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase
+class GroupByParityTests(
+ GroupByMixin,
+ PandasOnSparkTestUtils,
+ ReusedConnectTestCase,
):
pass
if __name__ == "__main__":
- from pyspark.pandas.tests.connect.test_parity_ops_on_diff_frames_groupby import * # noqa: F401
+ from pyspark.pandas.tests.connect.diff_frames_ops.test_parity_groupby import * # noqa
try:
import xmlrunner # type: ignore[import]
diff --git a/python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames_groupby.py b/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_groupby_shift.py
similarity index 79%
copy from python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames_groupby.py
copy to python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_groupby_shift.py
index 685ec5c45c5f..d1a116a273f8 100644
--- a/python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames_groupby.py
+++ b/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_groupby_shift.py
@@ -16,19 +16,21 @@
#
import unittest
-from pyspark.pandas.tests.test_ops_on_diff_frames_groupby import OpsOnDiffFramesGroupByTestsMixin
+from pyspark.pandas.tests.diff_frames_ops.test_groupby_shift import GroupByShiftMixin
from pyspark.testing.connectutils import ReusedConnectTestCase
from pyspark.testing.pandasutils import PandasOnSparkTestUtils
-class OpsOnDiffFramesGroupByParityTests(
- OpsOnDiffFramesGroupByTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase
+class GroupByShiftParityTests(
+ GroupByShiftMixin,
+ PandasOnSparkTestUtils,
+ ReusedConnectTestCase,
):
pass
if __name__ == "__main__":
- from pyspark.pandas.tests.connect.test_parity_ops_on_diff_frames_groupby import * # noqa: F401
+ from pyspark.pandas.tests.connect.diff_frames_ops.test_parity_groupby_shift import * # noqa
try:
import xmlrunner # type: ignore[import]
diff --git a/python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames_groupby.py b/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_groupby_split_apply_combine.py
similarity index 79%
copy from python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames_groupby.py
copy to python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_groupby_split_apply_combine.py
index 685ec5c45c5f..a4e32b144bf4 100644
--- a/python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames_groupby.py
+++ b/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_groupby_split_apply_combine.py
@@ -16,19 +16,21 @@
#
import unittest
-from pyspark.pandas.tests.test_ops_on_diff_frames_groupby import OpsOnDiffFramesGroupByTestsMixin
+from pyspark.pandas.tests.diff_frames_ops.test_groupby_split_apply_combine import GroupBySACMixin
from pyspark.testing.connectutils import ReusedConnectTestCase
from pyspark.testing.pandasutils import PandasOnSparkTestUtils
-class OpsOnDiffFramesGroupByParityTests(
- OpsOnDiffFramesGroupByTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase
+class GroupBySACParityTests(
+ GroupBySACMixin,
+ PandasOnSparkTestUtils,
+ ReusedConnectTestCase,
):
pass
if __name__ == "__main__":
- from pyspark.pandas.tests.connect.test_parity_ops_on_diff_frames_groupby import * # noqa: F401
+ from pyspark.pandas.tests.connect.diff_frames_ops.test_parity_groupby_split_apply_combine import * # noqa
try:
import xmlrunner # type: ignore[import]
diff --git a/python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames_groupby.py b/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_groupby_transform.py
similarity index 79%
rename from python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames_groupby.py
rename to python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_groupby_transform.py
index 685ec5c45c5f..64826d42563e 100644
--- a/python/pyspark/pandas/tests/connect/test_parity_ops_on_diff_frames_groupby.py
+++ b/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_groupby_transform.py
@@ -16,19 +16,21 @@
#
import unittest
-from pyspark.pandas.tests.test_ops_on_diff_frames_groupby import OpsOnDiffFramesGroupByTestsMixin
+from pyspark.pandas.tests.diff_frames_ops.test_groupby_transform import GroupByTransformMixin
from pyspark.testing.connectutils import ReusedConnectTestCase
from pyspark.testing.pandasutils import PandasOnSparkTestUtils
-class OpsOnDiffFramesGroupByParityTests(
- OpsOnDiffFramesGroupByTestsMixin, PandasOnSparkTestUtils, ReusedConnectTestCase
+class GroupByTransformParityTests(
+ GroupByTransformMixin,
+ PandasOnSparkTestUtils,
+ ReusedConnectTestCase,
):
pass
if __name__ == "__main__":
- from pyspark.pandas.tests.connect.test_parity_ops_on_diff_frames_groupby import * # noqa: F401
+ from pyspark.pandas.tests.connect.diff_frames_ops.test_parity_groupby_transform import * # noqa
try:
import xmlrunner # type: ignore[import]
diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby.py b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby.py
new file mode 100644
index 000000000000..25bff406bd9f
--- /dev/null
+++ b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby.py
@@ -0,0 +1,127 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import pandas as pd
+
+from pyspark import pandas as ps
+from pyspark.pandas.config import set_option, reset_option
+from pyspark.testing.pandasutils import PandasOnSparkTestCase
+from pyspark.testing.sqlutils import SQLTestUtils
+
+
+class GroupByMixin:
+ @classmethod
+ def setUpClass(cls):
+ super().setUpClass()
+ set_option("compute.ops_on_diff_frames", True)
+
+ @classmethod
+ def tearDownClass(cls):
+ reset_option("compute.ops_on_diff_frames")
+ super().tearDownClass()
+
+ def test_groupby_multiindex_columns(self):
+ pdf1 = pd.DataFrame(
+ {("y", "c"): [4, 2, 7, 3, None, 1, 1, 1, 2], ("z", "d"): list("abcdefght")}
+ )
+ pdf2 = pd.DataFrame(
+ {("x", "a"): [1, 2, 6, 4, 4, 6, 4, 3, 7], ("x", "b"): [4, 2, 7, 3, 3, 1, 1, 1, 2]}
+ )
+ psdf1 = ps.from_pandas(pdf1)
+ psdf2 = ps.from_pandas(pdf2)
+
+ self.assert_eq(
+ psdf1.groupby(psdf2[("x", "a")]).sum().sort_index(),
+ pdf1.groupby(pdf2[("x", "a")]).sum().sort_index(),
+ )
+
+ self.assert_eq(
+ psdf1.groupby(psdf2[("x", "a")], as_index=False)
+ .sum()
+ .sort_values(("y", "c"))
+ .reset_index(drop=True),
+ pdf1.groupby(pdf2[("x", "a")], as_index=False)
+ .sum()
+ .sort_values(("y", "c"))
+ .reset_index(drop=True),
+ )
+ self.assert_eq(
+ psdf1.groupby(psdf2[("x", "a")])[[("y", "c")]].sum().sort_index(),
+ pdf1.groupby(pdf2[("x", "a")])[[("y", "c")]].sum().sort_index(),
+ )
+
+ def test_duplicated_labels(self):
+ pdf1 = pd.DataFrame({"A": [3, 2, 1]})
+ pdf2 = pd.DataFrame({"A": [1, 2, 3]})
+ psdf1 = ps.from_pandas(pdf1)
+ psdf2 = ps.from_pandas(pdf2)
+
+ self.assert_eq(
+ psdf1.groupby(psdf2.A).sum().sort_index(), pdf1.groupby(pdf2.A).sum().sort_index()
+ )
+ self.assert_eq(
+ psdf1.groupby(psdf2.A, as_index=False).sum().sort_values("A").reset_index(drop=True),
+ pdf1.groupby(pdf2.A, as_index=False).sum().sort_values("A").reset_index(drop=True),
+ )
+
+ def test_head(self):
+ pdf = pd.DataFrame(
+ {
+ "a": [1, 1, 1, 1, 2, 2, 2, 3, 3, 3] * 3,
+ "b": [2, 3, 1, 4, 6, 9, 8, 10, 7, 5] * 3,
+ "c": [3, 5, 2, 5, 1, 2, 6, 4, 3, 6] * 3,
+ },
+ )
+ pkey = pd.Series([1, 1, 1, 1, 2, 2, 2, 3, 3, 3] * 3)
+ psdf = ps.from_pandas(pdf)
+ kkey = ps.from_pandas(pkey)
+
+ self.assert_eq(
+ pdf.groupby(pkey).head(2).sort_index(), psdf.groupby(kkey).head(2).sort_index()
+ )
+ self.assert_eq(
+ pdf.groupby("a")["b"].head(2).sort_index(), psdf.groupby("a")["b"].head(2).sort_index()
+ )
+ self.assert_eq(
+ pdf.groupby("a")[["b"]].head(2).sort_index(),
+ psdf.groupby("a")[["b"]].head(2).sort_index(),
+ )
+ self.assert_eq(
+ pdf.groupby([pkey, "b"]).head(2).sort_index(),
+ psdf.groupby([kkey, "b"]).head(2).sort_index(),
+ )
+
+
+class GroupByTests(
+ GroupByMixin,
+ PandasOnSparkTestCase,
+ SQLTestUtils,
+):
+ pass
+
+
+if __name__ == "__main__":
+ import unittest
+ from pyspark.pandas.tests.diff_frames_ops.test_groupby import * # noqa
+
+ try:
+ import xmlrunner
+
+ testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
+ except ImportError:
+ testRunner = None
+ unittest.main(testRunner=testRunner, verbosity=2)
diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_shift.py b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_shift.py
new file mode 100644
index 000000000000..260fba6cd2ac
--- /dev/null
+++ b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_shift.py
@@ -0,0 +1,85 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import pandas as pd
+
+from pyspark import pandas as ps
+from pyspark.pandas.config import set_option, reset_option
+from pyspark.testing.pandasutils import PandasOnSparkTestCase
+from pyspark.testing.sqlutils import SQLTestUtils
+
+
+class GroupByShiftMixin:
+ @classmethod
+ def setUpClass(cls):
+ super().setUpClass()
+ set_option("compute.ops_on_diff_frames", True)
+
+ @classmethod
+ def tearDownClass(cls):
+ reset_option("compute.ops_on_diff_frames")
+ super().tearDownClass()
+
+ def test_shift(self):
+ pdf = pd.DataFrame(
+ {
+ "a": [1, 1, 2, 2, 3, 3] * 3,
+ "b": [1, 1, 2, 2, 3, 4] * 3,
+ "c": [1, 4, 9, 16, 25, 36] * 3,
+ },
+ )
+ pkey = pd.Series([1, 1, 2, 2, 3, 4] * 3)
+ psdf = ps.from_pandas(pdf)
+ kkey = ps.from_pandas(pkey)
+
+ self.assert_eq(
+ psdf.groupby(kkey).shift().sort_index(), pdf.groupby(pkey).shift().sort_index()
+ )
+ self.assert_eq(
+ psdf.groupby(kkey)["a"].shift().sort_index(),
+ pdf.groupby(pkey)["a"].shift().sort_index(),
+ )
+ self.assert_eq(
+ psdf.groupby(kkey)[["a"]].shift().sort_index(),
+ pdf.groupby(pkey)[["a"]].shift().sort_index(),
+ )
+
+ self.assert_eq(
+ psdf.groupby(kkey).shift().sum(), pdf.groupby(pkey).shift().sum().astype(int)
+ )
+ self.assert_eq(psdf.groupby(kkey)["a"].shift().sum(), pdf.groupby(pkey)["a"].shift().sum())
+
+
+class GroupByShiftTests(
+ GroupByShiftMixin,
+ PandasOnSparkTestCase,
+ SQLTestUtils,
+):
+ pass
+
+
+if __name__ == "__main__":
+ import unittest
+ from pyspark.pandas.tests.diff_frames_ops.test_groupby_shift import * # noqa
+
+ try:
+ import xmlrunner
+
+ testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
+ except ImportError:
+ testRunner = None
+ unittest.main(testRunner=testRunner, verbosity=2)
diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_split_apply_combine.py b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_split_apply_combine.py
new file mode 100644
index 000000000000..0cfc30fa23b9
--- /dev/null
+++ b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_split_apply_combine.py
@@ -0,0 +1,113 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import pandas as pd
+
+from pyspark import pandas as ps
+from pyspark.pandas.config import set_option, reset_option
+from pyspark.testing.pandasutils import PandasOnSparkTestCase
+from pyspark.testing.sqlutils import SQLTestUtils
+
+
+class GroupBySACMixin:
+ @classmethod
+ def setUpClass(cls):
+ super().setUpClass()
+ set_option("compute.ops_on_diff_frames", True)
+
+ @classmethod
+ def tearDownClass(cls):
+ reset_option("compute.ops_on_diff_frames")
+ super().tearDownClass()
+
+ def test_split_apply_combine_on_series(self):
+ pdf1 = pd.DataFrame({"C": [0.362, 0.227, 1.267, -0.562], "B": [1, 2, 3, 4]})
+ pdf2 = pd.DataFrame({"A": [1, 1, 2, 2]})
+ psdf1 = ps.from_pandas(pdf1)
+ psdf2 = ps.from_pandas(pdf2)
+
+ for as_index in [True, False]:
+ if as_index:
+
+ def sort(df):
+ return df.sort_index()
+
+ else:
+
+ def sort(df):
+ return df.sort_values(list(df.columns)).reset_index(drop=True)
+
+ with self.subTest(as_index=as_index):
+ self.assert_eq(
+ sort(psdf1.groupby(psdf2.A, as_index=as_index).sum()),
+ sort(pdf1.groupby(pdf2.A, as_index=as_index).sum()),
+ )
+ self.assert_eq(
+ sort(psdf1.groupby(psdf2.A, as_index=as_index).B.sum()),
+ sort(pdf1.groupby(pdf2.A, as_index=as_index).B.sum()),
+ )
+ self.assert_eq(
+ sort(psdf1.groupby([psdf1.C, psdf2.A], as_index=as_index).sum()),
+ sort(pdf1.groupby([pdf1.C, pdf2.A], as_index=as_index).sum()),
+ )
+ self.assert_eq(
+ sort(psdf1.groupby([psdf1.C + 1, psdf2.A], as_index=as_index).sum()),
+ sort(pdf1.groupby([pdf1.C + 1, pdf2.A], as_index=as_index).sum()),
+ )
+
+ self.assert_eq(
+ psdf1.B.groupby(psdf2.A).sum().sort_index(),
+ pdf1.B.groupby(pdf2.A).sum().sort_index(),
+ )
+ self.assert_eq(
+ (psdf1.B + 1).groupby(psdf2.A).sum().sort_index(),
+ (pdf1.B + 1).groupby(pdf2.A).sum().sort_index(),
+ )
+
+ self.assert_eq(
+ psdf1.B.groupby(psdf2.A.rename()).sum().sort_index(),
+ pdf1.B.groupby(pdf2.A.rename()).sum().sort_index(),
+ )
+ self.assert_eq(
+ psdf1.B.rename().groupby(psdf2.A).sum().sort_index(),
+ pdf1.B.rename().groupby(pdf2.A).sum().sort_index(),
+ )
+ self.assert_eq(
+ psdf1.B.rename().groupby(psdf2.A.rename()).sum().sort_index(),
+ pdf1.B.rename().groupby(pdf2.A.rename()).sum().sort_index(),
+ )
+
+
+class GroupBySACTests(
+ GroupBySACMixin,
+ PandasOnSparkTestCase,
+ SQLTestUtils,
+):
+ pass
+
+
+if __name__ == "__main__":
+ import unittest
+ from pyspark.pandas.tests.diff_frames_ops.test_groupby_split_apply_combine import * # noqa
+
+ try:
+ import xmlrunner
+
+ testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
+ except ImportError:
+ testRunner = None
+ unittest.main(testRunner=testRunner, verbosity=2)
diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_transform.py b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_transform.py
new file mode 100644
index 000000000000..565096e41c83
--- /dev/null
+++ b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_transform.py
@@ -0,0 +1,82 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import pandas as pd
+
+from pyspark import pandas as ps
+from pyspark.pandas.config import set_option, reset_option
+from pyspark.testing.pandasutils import PandasOnSparkTestCase
+from pyspark.testing.sqlutils import SQLTestUtils
+
+
+class GroupByTransformMixin:
+ @classmethod
+ def setUpClass(cls):
+ super().setUpClass()
+ set_option("compute.ops_on_diff_frames", True)
+
+ @classmethod
+ def tearDownClass(cls):
+ reset_option("compute.ops_on_diff_frames")
+ super().tearDownClass()
+
+ def test_transform(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"],
+ )
+ pkey = pd.Series([1, 1, 2, 3, 5, 8])
+ psdf = ps.from_pandas(pdf)
+ kkey = ps.from_pandas(pkey)
+
+ self.assert_eq(
+ psdf.groupby(kkey).transform(lambda x: x + x.min()).sort_index(),
+ pdf.groupby(pkey).transform(lambda x: x + x.min()).sort_index(),
+ )
+ self.assert_eq(
+ psdf.groupby(kkey)["a"].transform(lambda x: x + x.min()).sort_index(),
+ pdf.groupby(pkey)["a"].transform(lambda x: x + x.min()).sort_index(),
+ )
+ self.assert_eq(
+ psdf.groupby(kkey)[["a"]].transform(lambda x: x + x.min()).sort_index(),
+ pdf.groupby(pkey)[["a"]].transform(lambda x: x + x.min()).sort_index(),
+ )
+ self.assert_eq(
+ psdf.groupby(["a", kkey]).transform(lambda x: x + x.min()).sort_index(),
+ pdf.groupby(["a", pkey]).transform(lambda x: x + x.min()).sort_index(),
+ )
+
+
+class GroupByTransformTests(
+ GroupByTransformMixin,
+ PandasOnSparkTestCase,
+ SQLTestUtils,
+):
+ pass
+
+
+if __name__ == "__main__":
+ import unittest
+ from pyspark.pandas.tests.diff_frames_ops.test_groupby_transform import * # noqa
+
+ try:
+ import xmlrunner
+
+ testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
+ except ImportError:
+ testRunner = None
+ unittest.main(testRunner=testRunner, verbosity=2)
diff --git a/python/pyspark/pandas/tests/test_ops_on_diff_frames_groupby.py b/python/pyspark/pandas/tests/test_ops_on_diff_frames_groupby.py
deleted file mode 100644
index 5c4c9283724c..000000000000
--- a/python/pyspark/pandas/tests/test_ops_on_diff_frames_groupby.py
+++ /dev/null
@@ -1,287 +0,0 @@
-#
-# Licensed to the Apache Software Foundation (ASF) under one or more
-# contributor license agreements. See the NOTICE file distributed with
-# this work for additional information regarding copyright ownership.
-# The ASF licenses this file to You under the Apache License, Version 2.0
-# (the "License"); you may not use this file except in compliance with
-# the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-
-import unittest
-
-import pandas as pd
-
-from pyspark import pandas as ps
-from pyspark.pandas.config import set_option, reset_option
-from pyspark.testing.pandasutils import PandasOnSparkTestCase
-from pyspark.testing.sqlutils import SQLTestUtils
-
-
-class OpsOnDiffFramesGroupByTestsMixin:
- @classmethod
- def setUpClass(cls):
- super().setUpClass()
- set_option("compute.ops_on_diff_frames", True)
-
- @classmethod
- def tearDownClass(cls):
- reset_option("compute.ops_on_diff_frames")
- super().tearDownClass()
-
- def test_groupby_multiindex_columns(self):
- pdf1 = pd.DataFrame(
- {("y", "c"): [4, 2, 7, 3, None, 1, 1, 1, 2], ("z", "d"): list("abcdefght")}
- )
- pdf2 = pd.DataFrame(
- {("x", "a"): [1, 2, 6, 4, 4, 6, 4, 3, 7], ("x", "b"): [4, 2, 7, 3, 3, 1, 1, 1, 2]}
- )
- psdf1 = ps.from_pandas(pdf1)
- psdf2 = ps.from_pandas(pdf2)
-
- self.assert_eq(
- psdf1.groupby(psdf2[("x", "a")]).sum().sort_index(),
- pdf1.groupby(pdf2[("x", "a")]).sum().sort_index(),
- )
-
- self.assert_eq(
- psdf1.groupby(psdf2[("x", "a")], as_index=False)
- .sum()
- .sort_values(("y", "c"))
- .reset_index(drop=True),
- pdf1.groupby(pdf2[("x", "a")], as_index=False)
- .sum()
- .sort_values(("y", "c"))
- .reset_index(drop=True),
- )
- self.assert_eq(
- psdf1.groupby(psdf2[("x", "a")])[[("y", "c")]].sum().sort_index(),
- pdf1.groupby(pdf2[("x", "a")])[[("y", "c")]].sum().sort_index(),
- )
-
- def test_split_apply_combine_on_series(self):
- pdf1 = pd.DataFrame({"C": [0.362, 0.227, 1.267, -0.562], "B": [1, 2, 3, 4]})
- pdf2 = pd.DataFrame({"A": [1, 1, 2, 2]})
- psdf1 = ps.from_pandas(pdf1)
- psdf2 = ps.from_pandas(pdf2)
-
- for as_index in [True, False]:
- if as_index:
-
- def sort(df):
- return df.sort_index()
-
- else:
-
- def sort(df):
- return df.sort_values(list(df.columns)).reset_index(drop=True)
-
- with self.subTest(as_index=as_index):
- self.assert_eq(
- sort(psdf1.groupby(psdf2.A, as_index=as_index).sum()),
- sort(pdf1.groupby(pdf2.A, as_index=as_index).sum()),
- )
- self.assert_eq(
- sort(psdf1.groupby(psdf2.A, as_index=as_index).B.sum()),
- sort(pdf1.groupby(pdf2.A, as_index=as_index).B.sum()),
- )
- self.assert_eq(
- sort(psdf1.groupby([psdf1.C, psdf2.A], as_index=as_index).sum()),
- sort(pdf1.groupby([pdf1.C, pdf2.A], as_index=as_index).sum()),
- )
- self.assert_eq(
- sort(psdf1.groupby([psdf1.C + 1, psdf2.A], as_index=as_index).sum()),
- sort(pdf1.groupby([pdf1.C + 1, pdf2.A], as_index=as_index).sum()),
- )
-
- self.assert_eq(
- psdf1.B.groupby(psdf2.A).sum().sort_index(),
- pdf1.B.groupby(pdf2.A).sum().sort_index(),
- )
- self.assert_eq(
- (psdf1.B + 1).groupby(psdf2.A).sum().sort_index(),
- (pdf1.B + 1).groupby(pdf2.A).sum().sort_index(),
- )
-
- self.assert_eq(
- psdf1.B.groupby(psdf2.A.rename()).sum().sort_index(),
- pdf1.B.groupby(pdf2.A.rename()).sum().sort_index(),
- )
- self.assert_eq(
- psdf1.B.rename().groupby(psdf2.A).sum().sort_index(),
- pdf1.B.rename().groupby(pdf2.A).sum().sort_index(),
- )
- self.assert_eq(
- psdf1.B.rename().groupby(psdf2.A.rename()).sum().sort_index(),
- pdf1.B.rename().groupby(pdf2.A.rename()).sum().sort_index(),
- )
-
- def test_duplicated_labels(self):
- pdf1 = pd.DataFrame({"A": [3, 2, 1]})
- pdf2 = pd.DataFrame({"A": [1, 2, 3]})
- psdf1 = ps.from_pandas(pdf1)
- psdf2 = ps.from_pandas(pdf2)
-
- self.assert_eq(
- psdf1.groupby(psdf2.A).sum().sort_index(), pdf1.groupby(pdf2.A).sum().sort_index()
- )
- self.assert_eq(
- psdf1.groupby(psdf2.A, as_index=False).sum().sort_values("A").reset_index(drop=True),
- pdf1.groupby(pdf2.A, as_index=False).sum().sort_values("A").reset_index(drop=True),
- )
-
- def test_transform(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"],
- )
- pkey = pd.Series([1, 1, 2, 3, 5, 8])
- psdf = ps.from_pandas(pdf)
- kkey = ps.from_pandas(pkey)
-
- self.assert_eq(
- psdf.groupby(kkey).transform(lambda x: x + x.min()).sort_index(),
- pdf.groupby(pkey).transform(lambda x: x + x.min()).sort_index(),
- )
- self.assert_eq(
- psdf.groupby(kkey)["a"].transform(lambda x: x + x.min()).sort_index(),
- pdf.groupby(pkey)["a"].transform(lambda x: x + x.min()).sort_index(),
- )
- self.assert_eq(
- psdf.groupby(kkey)[["a"]].transform(lambda x: x + x.min()).sort_index(),
- pdf.groupby(pkey)[["a"]].transform(lambda x: x + x.min()).sort_index(),
- )
- self.assert_eq(
- psdf.groupby(["a", kkey]).transform(lambda x: x + x.min()).sort_index(),
- pdf.groupby(["a", pkey]).transform(lambda x: x + x.min()).sort_index(),
- )
-
- def test_head(self):
- pdf = pd.DataFrame(
- {
- "a": [1, 1, 1, 1, 2, 2, 2, 3, 3, 3] * 3,
- "b": [2, 3, 1, 4, 6, 9, 8, 10, 7, 5] * 3,
- "c": [3, 5, 2, 5, 1, 2, 6, 4, 3, 6] * 3,
- },
- )
- pkey = pd.Series([1, 1, 1, 1, 2, 2, 2, 3, 3, 3] * 3)
- psdf = ps.from_pandas(pdf)
- kkey = ps.from_pandas(pkey)
-
- self.assert_eq(
- pdf.groupby(pkey).head(2).sort_index(), psdf.groupby(kkey).head(2).sort_index()
- )
- self.assert_eq(
- pdf.groupby("a")["b"].head(2).sort_index(), psdf.groupby("a")["b"].head(2).sort_index()
- )
- self.assert_eq(
- pdf.groupby("a")[["b"]].head(2).sort_index(),
- psdf.groupby("a")[["b"]].head(2).sort_index(),
- )
- self.assert_eq(
- pdf.groupby([pkey, "b"]).head(2).sort_index(),
- psdf.groupby([kkey, "b"]).head(2).sort_index(),
- )
-
- def test_shift(self):
- pdf = pd.DataFrame(
- {
- "a": [1, 1, 2, 2, 3, 3] * 3,
- "b": [1, 1, 2, 2, 3, 4] * 3,
- "c": [1, 4, 9, 16, 25, 36] * 3,
- },
- )
- pkey = pd.Series([1, 1, 2, 2, 3, 4] * 3)
- psdf = ps.from_pandas(pdf)
- kkey = ps.from_pandas(pkey)
-
- self.assert_eq(
- psdf.groupby(kkey).shift().sort_index(), pdf.groupby(pkey).shift().sort_index()
- )
- self.assert_eq(
- psdf.groupby(kkey)["a"].shift().sort_index(),
- pdf.groupby(pkey)["a"].shift().sort_index(),
- )
- self.assert_eq(
- psdf.groupby(kkey)[["a"]].shift().sort_index(),
- pdf.groupby(pkey)[["a"]].shift().sort_index(),
- )
-
- self.assert_eq(
- psdf.groupby(kkey).shift().sum(), pdf.groupby(pkey).shift().sum().astype(int)
- )
- self.assert_eq(psdf.groupby(kkey)["a"].shift().sum(), pdf.groupby(pkey)["a"].shift().sum())
-
- def test_fillna(self):
- pdf = pd.DataFrame(
- {
- "A": [1, 1, 2, 2] * 3,
- "B": [2, 4, None, 3] * 3,
- "C": [None, None, None, 1] * 3,
- "D": [0, 1, 5, 4] * 3,
- }
- )
- pkey = pd.Series([1, 1, 2, 2] * 3)
- psdf = ps.from_pandas(pdf)
- kkey = ps.from_pandas(pkey)
-
- self.assert_eq(
- psdf.groupby(kkey).fillna(0).sort_index(), pdf.groupby(pkey).fillna(0).sort_index()
- )
- self.assert_eq(
- psdf.groupby(kkey)["C"].fillna(0).sort_index(),
- pdf.groupby(pkey)["C"].fillna(0).sort_index(),
- )
- self.assert_eq(
- psdf.groupby(kkey)[["C"]].fillna(0).sort_index(),
- pdf.groupby(pkey)[["C"]].fillna(0).sort_index(),
- )
- self.assert_eq(
- psdf.groupby(kkey).fillna(method="bfill").sort_index(),
- pdf.groupby(pkey).fillna(method="bfill").sort_index(),
- )
- self.assert_eq(
- psdf.groupby(kkey)["C"].fillna(method="bfill").sort_index(),
- pdf.groupby(pkey)["C"].fillna(method="bfill").sort_index(),
- )
- self.assert_eq(
- psdf.groupby(kkey)[["C"]].fillna(method="bfill").sort_index(),
- pdf.groupby(pkey)[["C"]].fillna(method="bfill").sort_index(),
- )
- self.assert_eq(
- psdf.groupby(kkey).fillna(method="ffill").sort_index(),
- pdf.groupby(pkey).fillna(method="ffill").sort_index(),
- )
- self.assert_eq(
- psdf.groupby(kkey)["C"].fillna(method="ffill").sort_index(),
- pdf.groupby(pkey)["C"].fillna(method="ffill").sort_index(),
- )
- self.assert_eq(
- psdf.groupby(kkey)[["C"]].fillna(method="ffill").sort_index(),
- pdf.groupby(pkey)[["C"]].fillna(method="ffill").sort_index(),
- )
-
-
-class OpsOnDiffFramesGroupByTests(
- OpsOnDiffFramesGroupByTestsMixin, PandasOnSparkTestCase, SQLTestUtils
-):
- pass
-
-
-if __name__ == "__main__":
- from pyspark.pandas.tests.test_ops_on_diff_frames_groupby import * # noqa: F401
-
- try:
- import xmlrunner
-
- testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
- except ImportError:
- testRunner = None
- unittest.main(testRunner=testRunner, verbosity=2)
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org