This is an automated email from the ASF dual-hosted git repository.
HyukjinKwon pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 89b66fe831f1 [SPARK-58058][PYTHON][TEST] Split slow PySpark test
suites into smaller modules
89b66fe831f1 is described below
commit 89b66fe831f1f3a82dd7c9b629cb0197297203cc
Author: Hyukjin Kwon <[email protected]>
AuthorDate: Fri Jul 10 06:45:45 2026 +0900
[SPARK-58058][PYTHON][TEST] Split slow PySpark test suites into smaller
modules
### What changes were proposed in this pull request?
This PR splits several large/slow PySpark test suites into smaller modules
so
they run as separate test targets and parallelize better, reducing
wall-clock
time. No test logic is changed; test methods are only relocated, following
the
existing convention already used in the repo (e.g. `test_split_apply_*`,
`test_stat_*`).
The following suites are split, each with a matching Spark Connect parity
module:
- `pyspark.sql.tests.pandas.test_pandas_cogrouped_map`: move the heavier,
self-contained tests (`test_with_window_function`, `test_with_local_data`,
`test_arrow_batch_slicing`, `test_cogroup_apply_in_pandas_with_logging`)
into
a new `test_pandas_cogrouped_map_misc` module.
- `pyspark.sql.tests.arrow.test_arrow_cogrouped_map`: extract shared helpers
into `CogroupedMapInArrowTestsFuncMixin` and move the heavier tests into
a new
`test_arrow_cogrouped_map_misc` module.
- `pyspark.pandas.tests.data_type_ops.test_num_ops`: move
`IntegralExtensionOpsTestsMixin` and `FractionalExtensionOpsTestsMixin`
into
their own modules.
- `pyspark.pandas.tests.groupby.test_stat`: move `test_median` into a new
`test_stat_median` module.
- `pyspark.pandas.tests.groupby.test_split_apply`: move the `mean` case
into a
new `test_split_apply_mean` module.
New modules are registered in `dev/sparktestsupport/modules.py`.
### Why are the changes needed?
These suites are slow in CI. Splitting them lets the test runner schedule
the
pieces in parallel, lowering total wall-clock time without dropping
coverage.
### Does this PR introduce _any_ user-facing change?
No. Test-only change.
### How was this patch tested?
Existing tests, only relocated across modules. Verified that the total
number
of test methods is unchanged for each split suite, that all new and modified
modules import and collect, and that `flake8` passes.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Cursor
Closes #57155 from HyukjinKwon/hyukjin/split-slow-pyspark-tests.
Authored-by: Hyukjin Kwon <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>
(cherry picked from commit 79f6dfa4b99865f581d9c073bddd352e666bcb47)
Signed-off-by: Hyukjin Kwon <[email protected]>
---
dev/sparktestsupport/modules.py | 16 ++
.../connect/data_type_ops/test_parity_num_ops.py | 24 +-
...ps.py => test_parity_num_ops_fractional_ext.py} | 22 +-
...st_parity_num_ops_fractional_ext_astype_cmp.py} | 28 +--
..._ops.py => test_parity_num_ops_integral_ext.py} | 22 +-
...test_parity_num_ops_integral_ext_astype_cmp.py} | 28 +--
.../test_parity_split_apply_mean.py} | 32 +--
.../test_parity_stat_median.py} | 32 +--
.../pandas/tests/data_type_ops/test_num_ops.py | 254 ---------------------
.../data_type_ops/test_num_ops_fractional_ext.py | 84 +++++++
.../test_num_ops_fractional_ext_astype_cmp.py | 120 ++++++++++
.../data_type_ops/test_num_ops_integral_ext.py | 106 +++++++++
.../test_num_ops_integral_ext_astype_cmp.py | 117 ++++++++++
.../pandas/tests/groupby/test_split_apply.py | 1 -
.../test_split_apply_mean.py} | 40 +---
python/pyspark/pandas/tests/groupby/test_stat.py | 31 ---
.../pandas/tests/groupby/test_stat_median.py | 84 +++++++
.../sql/tests/arrow/test_arrow_cogrouped_map.py | 193 +---------------
.../tests/arrow/test_arrow_cogrouped_map_misc.py | 240 +++++++++++++++++++
.../arrow/test_parity_arrow_cogrouped_map_misc.py} | 34 +--
.../test_parity_pandas_cogrouped_map_misc.py} | 32 +--
.../sql/tests/pandas/test_pandas_cogrouped_map.py | 189 ---------------
.../tests/pandas/test_pandas_cogrouped_map_misc.py | 239 +++++++++++++++++++
23 files changed, 1052 insertions(+), 916 deletions(-)
diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py
index 48e9fab4aa2b..1e1cd1e3af73 100644
--- a/dev/sparktestsupport/modules.py
+++ b/dev/sparktestsupport/modules.py
@@ -591,6 +591,7 @@ pyspark_sql = Module(
"pyspark.sql.tests.arrow.test_arrow",
"pyspark.sql.tests.arrow.test_arrow_map",
"pyspark.sql.tests.arrow.test_arrow_cogrouped_map",
+ "pyspark.sql.tests.arrow.test_arrow_cogrouped_map_misc",
"pyspark.sql.tests.arrow.test_arrow_grouped_map",
"pyspark.sql.tests.arrow.test_arrow_python_udf",
"pyspark.sql.tests.arrow.test_arrow_udf",
@@ -600,6 +601,7 @@ pyspark_sql = Module(
"pyspark.sql.tests.arrow.test_arrow_udf_typehints",
"pyspark.sql.tests.arrow.test_arrow_udtf",
"pyspark.sql.tests.pandas.test_pandas_cogrouped_map",
+ "pyspark.sql.tests.pandas.test_pandas_cogrouped_map_misc",
"pyspark.sql.tests.pandas.test_pandas_grouped_map",
"pyspark.sql.tests.pandas.test_pandas_map",
"pyspark.sql.tests.pandas.test_pandas_udf",
@@ -911,6 +913,10 @@ pyspark_pandas = Module(
"pyspark.pandas.tests.data_type_ops.test_datetime_ops",
"pyspark.pandas.tests.data_type_ops.test_null_ops",
"pyspark.pandas.tests.data_type_ops.test_num_ops",
+ "pyspark.pandas.tests.data_type_ops.test_num_ops_integral_ext",
+
"pyspark.pandas.tests.data_type_ops.test_num_ops_integral_ext_astype_cmp",
+ "pyspark.pandas.tests.data_type_ops.test_num_ops_fractional_ext",
+
"pyspark.pandas.tests.data_type_ops.test_num_ops_fractional_ext_astype_cmp",
"pyspark.pandas.tests.data_type_ops.test_num_arithmetic",
"pyspark.pandas.tests.data_type_ops.test_num_mod",
"pyspark.pandas.tests.data_type_ops.test_num_mul_div",
@@ -1070,6 +1076,7 @@ pyspark_pandas_slow = Module(
"pyspark.pandas.tests.groupby.test_size",
"pyspark.pandas.tests.groupby.test_split_apply",
"pyspark.pandas.tests.groupby.test_split_apply_count",
+ "pyspark.pandas.tests.groupby.test_split_apply_mean",
"pyspark.pandas.tests.groupby.test_split_apply_first",
"pyspark.pandas.tests.groupby.test_split_apply_last",
"pyspark.pandas.tests.groupby.test_split_apply_min_max",
@@ -1080,6 +1087,7 @@ pyspark_pandas_slow = Module(
"pyspark.pandas.tests.groupby.test_stat_adv",
"pyspark.pandas.tests.groupby.test_stat_ddof",
"pyspark.pandas.tests.groupby.test_stat_func",
+ "pyspark.pandas.tests.groupby.test_stat_median",
"pyspark.pandas.tests.groupby.test_stat_prod",
"pyspark.pandas.tests.groupby.test_value_counts",
"pyspark.pandas.tests.diff_frames_ops.test_align",
@@ -1217,6 +1225,7 @@ pyspark_connect = Module(
"pyspark.sql.tests.connect.arrow.test_parity_arrow_map",
"pyspark.sql.tests.connect.arrow.test_parity_arrow_grouped_map",
"pyspark.sql.tests.connect.arrow.test_parity_arrow_cogrouped_map",
+ "pyspark.sql.tests.connect.arrow.test_parity_arrow_cogrouped_map_misc",
"pyspark.sql.tests.connect.arrow.test_parity_arrow_python_udf",
"pyspark.sql.tests.connect.arrow.test_parity_arrow_udf",
"pyspark.sql.tests.connect.arrow.test_parity_arrow_udf_scalar",
@@ -1226,6 +1235,7 @@ pyspark_connect = Module(
"pyspark.sql.tests.connect.pandas.test_parity_pandas_map",
"pyspark.sql.tests.connect.pandas.test_parity_pandas_grouped_map",
"pyspark.sql.tests.connect.pandas.test_parity_pandas_cogrouped_map",
+
"pyspark.sql.tests.connect.pandas.test_parity_pandas_cogrouped_map_misc",
"pyspark.sql.tests.connect.pandas.test_parity_pandas_udf",
"pyspark.sql.tests.connect.pandas.test_parity_pandas_udf_scalar",
"pyspark.sql.tests.connect.pandas.test_parity_pandas_udf_grouped_agg",
@@ -1346,6 +1356,10 @@ pyspark_pandas_connect = Module(
"pyspark.pandas.tests.connect.data_type_ops.test_parity_datetime_ops",
"pyspark.pandas.tests.connect.data_type_ops.test_parity_null_ops",
"pyspark.pandas.tests.connect.data_type_ops.test_parity_num_ops",
+
"pyspark.pandas.tests.connect.data_type_ops.test_parity_num_ops_integral_ext",
+
"pyspark.pandas.tests.connect.data_type_ops.test_parity_num_ops_integral_ext_astype_cmp",
+
"pyspark.pandas.tests.connect.data_type_ops.test_parity_num_ops_fractional_ext",
+
"pyspark.pandas.tests.connect.data_type_ops.test_parity_num_ops_fractional_ext_astype_cmp",
"pyspark.pandas.tests.connect.data_type_ops.test_parity_num_reverse",
"pyspark.pandas.tests.connect.data_type_ops.test_parity_string_ops",
"pyspark.pandas.tests.connect.data_type_ops.test_parity_udt_ops",
@@ -1489,6 +1503,7 @@ pyspark_pandas_slow_connect = Module(
"pyspark.pandas.tests.connect.groupby.test_parity_stat_adv",
"pyspark.pandas.tests.connect.groupby.test_parity_stat_ddof",
"pyspark.pandas.tests.connect.groupby.test_parity_stat_func",
+ "pyspark.pandas.tests.connect.groupby.test_parity_stat_median",
"pyspark.pandas.tests.connect.groupby.test_parity_stat_prod",
"pyspark.pandas.tests.connect.groupby.test_parity_aggregate",
"pyspark.pandas.tests.connect.groupby.test_parity_apply_func",
@@ -1499,6 +1514,7 @@ pyspark_pandas_slow_connect = Module(
"pyspark.pandas.tests.connect.groupby.test_parity_split_apply",
"pyspark.pandas.tests.connect.groupby.test_parity_split_apply_count",
"pyspark.pandas.tests.connect.groupby.test_parity_split_apply_first",
+ "pyspark.pandas.tests.connect.groupby.test_parity_split_apply_mean",
"pyspark.pandas.tests.connect.groupby.test_parity_split_apply_last",
"pyspark.pandas.tests.connect.groupby.test_parity_split_apply_min_max",
"pyspark.pandas.tests.connect.groupby.test_parity_split_apply_skew",
diff --git
a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
index ed8dd1083c48..faeef9b0dd4a 100644
--- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
+++ b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
@@ -15,11 +15,7 @@
# limitations under the License.
#
-from pyspark.pandas.tests.data_type_ops.test_num_ops import (
- NumOpsTestsMixin,
- IntegralExtensionOpsTestsMixin,
- FractionalExtensionOpsTestsMixin,
-)
+from pyspark.pandas.tests.data_type_ops.test_num_ops import NumOpsTestsMixin
from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase
from pyspark.testing.pandasutils import PandasOnSparkTestUtils
from pyspark.testing.connectutils import ReusedConnectTestCase
@@ -34,24 +30,6 @@ class NumOpsParityTests(
pass
-class IntegralExtensionOpsParityTests(
- IntegralExtensionOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
- ReusedConnectTestCase,
-):
- pass
-
-
-class FractionalExtensionOpsParityTests(
- FractionalExtensionOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
- ReusedConnectTestCase,
-):
- pass
-
-
if __name__ == "__main__":
from pyspark.testing import main
diff --git
a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_fractional_ext.py
similarity index 75%
copy from
python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
copy to
python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_fractional_ext.py
index ed8dd1083c48..da14d21cc3ed 100644
--- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
+++
b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_fractional_ext.py
@@ -15,9 +15,7 @@
# limitations under the License.
#
-from pyspark.pandas.tests.data_type_ops.test_num_ops import (
- NumOpsTestsMixin,
- IntegralExtensionOpsTestsMixin,
+from pyspark.pandas.tests.data_type_ops.test_num_ops_fractional_ext import (
FractionalExtensionOpsTestsMixin,
)
from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase
@@ -25,24 +23,6 @@ from pyspark.testing.pandasutils import
PandasOnSparkTestUtils
from pyspark.testing.connectutils import ReusedConnectTestCase
-class NumOpsParityTests(
- NumOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
- ReusedConnectTestCase,
-):
- pass
-
-
-class IntegralExtensionOpsParityTests(
- IntegralExtensionOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
- ReusedConnectTestCase,
-):
- pass
-
-
class FractionalExtensionOpsParityTests(
FractionalExtensionOpsTestsMixin,
PandasOnSparkTestUtils,
diff --git
a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_fractional_ext_astype_cmp.py
similarity index 68%
copy from
python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
copy to
python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_fractional_ext_astype_cmp.py
index ed8dd1083c48..f140db0a13f2 100644
--- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
+++
b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_fractional_ext_astype_cmp.py
@@ -15,36 +15,16 @@
# limitations under the License.
#
-from pyspark.pandas.tests.data_type_ops.test_num_ops import (
- NumOpsTestsMixin,
- IntegralExtensionOpsTestsMixin,
- FractionalExtensionOpsTestsMixin,
+from pyspark.pandas.tests.data_type_ops.test_num_ops_fractional_ext_astype_cmp
import (
+ FractionalExtensionAstypeCmpOpsTestsMixin,
)
from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase
from pyspark.testing.pandasutils import PandasOnSparkTestUtils
from pyspark.testing.connectutils import ReusedConnectTestCase
-class NumOpsParityTests(
- NumOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
- ReusedConnectTestCase,
-):
- pass
-
-
-class IntegralExtensionOpsParityTests(
- IntegralExtensionOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
- ReusedConnectTestCase,
-):
- pass
-
-
-class FractionalExtensionOpsParityTests(
- FractionalExtensionOpsTestsMixin,
+class FractionalExtensionAstypeCmpOpsParityTests(
+ FractionalExtensionAstypeCmpOpsTestsMixin,
PandasOnSparkTestUtils,
OpsTestBase,
ReusedConnectTestCase,
diff --git
a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_integral_ext.py
similarity index 75%
copy from
python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
copy to
python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_integral_ext.py
index ed8dd1083c48..863b09d1a737 100644
--- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
+++
b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_integral_ext.py
@@ -15,25 +15,14 @@
# limitations under the License.
#
-from pyspark.pandas.tests.data_type_ops.test_num_ops import (
- NumOpsTestsMixin,
+from pyspark.pandas.tests.data_type_ops.test_num_ops_integral_ext import (
IntegralExtensionOpsTestsMixin,
- FractionalExtensionOpsTestsMixin,
)
from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase
from pyspark.testing.pandasutils import PandasOnSparkTestUtils
from pyspark.testing.connectutils import ReusedConnectTestCase
-class NumOpsParityTests(
- NumOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
- ReusedConnectTestCase,
-):
- pass
-
-
class IntegralExtensionOpsParityTests(
IntegralExtensionOpsTestsMixin,
PandasOnSparkTestUtils,
@@ -43,15 +32,6 @@ class IntegralExtensionOpsParityTests(
pass
-class FractionalExtensionOpsParityTests(
- FractionalExtensionOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
- ReusedConnectTestCase,
-):
- pass
-
-
if __name__ == "__main__":
from pyspark.testing import main
diff --git
a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_integral_ext_astype_cmp.py
similarity index 68%
copy from
python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
copy to
python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_integral_ext_astype_cmp.py
index ed8dd1083c48..af85dc5c8fca 100644
--- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
+++
b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_integral_ext_astype_cmp.py
@@ -15,36 +15,16 @@
# limitations under the License.
#
-from pyspark.pandas.tests.data_type_ops.test_num_ops import (
- NumOpsTestsMixin,
- IntegralExtensionOpsTestsMixin,
- FractionalExtensionOpsTestsMixin,
+from pyspark.pandas.tests.data_type_ops.test_num_ops_integral_ext_astype_cmp
import (
+ IntegralExtensionAstypeCmpOpsTestsMixin,
)
from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase
from pyspark.testing.pandasutils import PandasOnSparkTestUtils
from pyspark.testing.connectutils import ReusedConnectTestCase
-class NumOpsParityTests(
- NumOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
- ReusedConnectTestCase,
-):
- pass
-
-
-class IntegralExtensionOpsParityTests(
- IntegralExtensionOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
- ReusedConnectTestCase,
-):
- pass
-
-
-class FractionalExtensionOpsParityTests(
- FractionalExtensionOpsTestsMixin,
+class IntegralExtensionAstypeCmpOpsParityTests(
+ IntegralExtensionAstypeCmpOpsTestsMixin,
PandasOnSparkTestUtils,
OpsTestBase,
ReusedConnectTestCase,
diff --git
a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
b/python/pyspark/pandas/tests/connect/groupby/test_parity_split_apply_mean.py
similarity index 62%
copy from
python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
copy to
python/pyspark/pandas/tests/connect/groupby/test_parity_split_apply_mean.py
index ed8dd1083c48..d02c4ef5a16e 100644
--- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
+++
b/python/pyspark/pandas/tests/connect/groupby/test_parity_split_apply_mean.py
@@ -15,38 +15,14 @@
# limitations under the License.
#
-from pyspark.pandas.tests.data_type_ops.test_num_ops import (
- NumOpsTestsMixin,
- IntegralExtensionOpsTestsMixin,
- FractionalExtensionOpsTestsMixin,
-)
-from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase
-from pyspark.testing.pandasutils import PandasOnSparkTestUtils
+from pyspark.pandas.tests.groupby.test_split_apply_mean import
GroupbySplitApplyMeanMixin
from pyspark.testing.connectutils import ReusedConnectTestCase
+from pyspark.testing.pandasutils import PandasOnSparkTestUtils
-class NumOpsParityTests(
- NumOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
- ReusedConnectTestCase,
-):
- pass
-
-
-class IntegralExtensionOpsParityTests(
- IntegralExtensionOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
- ReusedConnectTestCase,
-):
- pass
-
-
-class FractionalExtensionOpsParityTests(
- FractionalExtensionOpsTestsMixin,
+class GroupbySplitApplyMeanParityTests(
+ GroupbySplitApplyMeanMixin,
PandasOnSparkTestUtils,
- OpsTestBase,
ReusedConnectTestCase,
):
pass
diff --git
a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
b/python/pyspark/pandas/tests/connect/groupby/test_parity_stat_median.py
similarity index 62%
copy from
python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
copy to python/pyspark/pandas/tests/connect/groupby/test_parity_stat_median.py
index ed8dd1083c48..a45574690f5d 100644
--- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
+++ b/python/pyspark/pandas/tests/connect/groupby/test_parity_stat_median.py
@@ -15,38 +15,14 @@
# limitations under the License.
#
-from pyspark.pandas.tests.data_type_ops.test_num_ops import (
- NumOpsTestsMixin,
- IntegralExtensionOpsTestsMixin,
- FractionalExtensionOpsTestsMixin,
-)
-from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase
-from pyspark.testing.pandasutils import PandasOnSparkTestUtils
+from pyspark.pandas.tests.groupby.test_stat_median import
GroupbyStatMedianMixin
from pyspark.testing.connectutils import ReusedConnectTestCase
+from pyspark.testing.pandasutils import PandasOnSparkTestUtils
-class NumOpsParityTests(
- NumOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
- ReusedConnectTestCase,
-):
- pass
-
-
-class IntegralExtensionOpsParityTests(
- IntegralExtensionOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
- ReusedConnectTestCase,
-):
- pass
-
-
-class FractionalExtensionOpsParityTests(
- FractionalExtensionOpsTestsMixin,
+class GroupbyStatMedianParityTests(
+ GroupbyStatMedianMixin,
PandasOnSparkTestUtils,
- OpsTestBase,
ReusedConnectTestCase,
):
pass
diff --git a/python/pyspark/pandas/tests/data_type_ops/test_num_ops.py
b/python/pyspark/pandas/tests/data_type_ops/test_num_ops.py
index 9088262308c2..4e95ee1cb0aa 100644
--- a/python/pyspark/pandas/tests/data_type_ops/test_num_ops.py
+++ b/python/pyspark/pandas/tests/data_type_ops/test_num_ops.py
@@ -15,8 +15,6 @@
# limitations under the License.
#
-import unittest
-
import pandas as pd
import numpy as np
@@ -25,10 +23,6 @@ from pyspark.pandas.config import option_context
from pyspark.testing.pandasutils import PandasOnSparkTestCase
from pyspark.testing.utils import is_ansi_mode_test
from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase
-from pyspark.pandas.typedef.typehints import (
- extension_dtypes_available,
- extension_float_dtypes_available,
-)
from pyspark.sql.types import DecimalType, IntegralType
@@ -199,238 +193,6 @@ class NumOpsTestsMixin:
self.assert_eq(pser >= pser, psser >= psser)
[email protected](not extension_dtypes_available, "pandas extension dtypes are
not available")
-class IntegralExtensionOpsTestsMixin:
- @property
- def intergral_extension_psers(self):
- return [pd.Series([1, 2, 3, None], dtype=dtype) for dtype in
self.integral_extension_dtypes]
-
- @property
- def intergral_extension_pssers(self):
- return [ps.from_pandas(pser) for pser in
self.intergral_extension_psers]
-
- @property
- def intergral_extension_pser_psser_pairs(self):
- return zip(self.intergral_extension_psers,
self.intergral_extension_pssers)
-
- def test_from_to_pandas(self):
- for pser, psser in self.intergral_extension_pser_psser_pairs:
- self.check_extension(pser, psser._to_pandas())
- self.check_extension(ps.from_pandas(pser), psser)
-
- def test_isnull(self):
- for pser, psser in self.intergral_extension_pser_psser_pairs:
- self.assert_eq(pser.isnull(), psser.isnull())
-
- def test_astype(self):
- for pser, psser in self.intergral_extension_pser_psser_pairs:
- for dtype in self.extension_dtypes:
- if dtype in self.string_extension_dtype:
- self.check_extension(pser.astype(dtype),
psser.astype(dtype))
- else:
- self.check_extension(pser.astype(dtype),
psser.astype(dtype))
- for pser, psser in self.intergral_extension_pser_psser_pairs:
- self.assert_eq(pser.astype(float), psser.astype(float))
- self.assert_eq(pser.astype(np.float32), psser.astype(np.float32))
- with ps.option_context("compute.eager_check", True):
- self.assertRaisesRegex(
- ValueError,
- "Cannot convert integrals with missing values to bool",
- lambda: psser.astype(bool),
- )
- self.assertRaisesRegex(
- ValueError,
- "Cannot convert integrals with missing values to integer",
- lambda: psser.astype(int),
- )
- self.assertRaisesRegex(
- ValueError,
- "Cannot convert integrals with missing values to integer",
- lambda: psser.astype(np.int32),
- )
- with ps.option_context("compute.eager_check", False):
- psser.astype(bool)
- psser.astype(int)
- psser.astype(np.int32)
-
- def test_neg(self):
- for pser, psser in self.intergral_extension_pser_psser_pairs:
- self.check_extension(-pser, -psser)
-
- def test_abs(self):
- for pser, psser in self.intergral_extension_pser_psser_pairs:
- self.check_extension(abs(pser), abs(psser))
-
- def test_invert(self):
- for pser, psser in self.intergral_extension_pser_psser_pairs:
- self.check_extension(~pser, ~psser)
-
- def test_eq(self):
- with option_context("compute.ops_on_diff_frames", True):
- for pser, psser in self.intergral_extension_pser_psser_pairs:
- self.check_extension(pser == pser, (psser ==
psser).sort_index())
-
- def test_ne(self):
- with option_context("compute.ops_on_diff_frames", True):
- for pser, psser in self.intergral_extension_pser_psser_pairs:
- self.check_extension(pser != pser, (psser !=
psser).sort_index())
-
- def test_lt(self):
- with option_context("compute.ops_on_diff_frames", True):
- for pser, psser in self.intergral_extension_pser_psser_pairs:
- self.check_extension(pser < pser, (psser < psser).sort_index())
-
- def test_le(self):
- with option_context("compute.ops_on_diff_frames", True):
- for pser, psser in self.intergral_extension_pser_psser_pairs:
- self.check_extension(pser <= pser, (psser <=
psser).sort_index())
-
- def test_gt(self):
- with option_context("compute.ops_on_diff_frames", True):
- for pser, psser in self.intergral_extension_pser_psser_pairs:
- self.check_extension(pser > pser, (psser > psser).sort_index())
-
- def test_ge(self):
- with option_context("compute.ops_on_diff_frames", True):
- for pser, psser in self.intergral_extension_pser_psser_pairs:
- self.check_extension(pser >= pser, (psser >=
psser).sort_index())
-
- def test_xor(self):
- for psser in self.intergral_extension_pssers:
- self.assertRaisesRegex(
- TypeError,
- "XOR can not be applied to given types.",
- lambda: psser ^ 1,
- )
- self.assertRaisesRegex(
- TypeError,
- "XOR can not be applied to given types.",
- lambda: psser ^ psser,
- )
- self.assertRaisesRegex(
- TypeError,
- "XOR can not be applied to given types.",
- lambda: psser ^ False,
- )
-
- def test_rxor(self):
- for psser in self.intergral_extension_pssers:
- self.assertRaisesRegex(
- TypeError,
- "XOR can not be applied to given types.",
- lambda: 1 ^ psser,
- )
- self.assertRaisesRegex(
- TypeError,
- "XOR can not be applied to given types.",
- lambda: False ^ psser,
- )
-
-
[email protected](
- not extension_float_dtypes_available, "pandas extension float dtypes are
not available"
-)
-class FractionalExtensionOpsTestsMixin:
- @property
- def fractional_extension_psers(self):
- return [
- pd.Series([0.1, 0.2, 0.3, None], dtype=dtype)
- for dtype in self.fractional_extension_dtypes
- ]
-
- @property
- def fractional_extension_pssers(self):
- return [ps.from_pandas(pser) for pser in
self.fractional_extension_psers]
-
- @property
- def fractional_extension_pser_psser_pairs(self):
- return zip(self.fractional_extension_psers,
self.fractional_extension_pssers)
-
- def test_from_to_pandas(self):
- for pser, psser in self.fractional_extension_pser_psser_pairs:
- self.check_extension(pser, psser._to_pandas())
- self.check_extension(ps.from_pandas(pser), psser)
-
- def test_isnull(self):
- for pser, psser in self.fractional_extension_pser_psser_pairs:
- self.assert_eq(pser.isnull(), psser.isnull())
-
- def test_astype(self):
- for pser, psser in self.fractional_extension_pser_psser_pairs:
- for dtype in self.extension_dtypes:
- self.check_extension(pser.astype(dtype), psser.astype(dtype))
- for pser, psser in self.fractional_extension_pser_psser_pairs:
- self.assert_eq(pser.astype(float), psser.astype(float))
- self.assert_eq(pser.astype("category"), psser.astype("category"))
- self.assert_eq(pser.astype(np.float32), psser.astype(np.float32))
- with ps.option_context("compute.eager_check", True):
- self.assertRaisesRegex(
- ValueError,
- "Cannot convert fractions with missing values to bool",
- lambda: psser.astype(bool),
- )
- self.assertRaisesRegex(
- ValueError,
- "Cannot convert fractions with missing values to integer",
- lambda: psser.astype(int),
- )
- self.assertRaisesRegex(
- ValueError,
- "Cannot convert fractions with missing values to integer",
- lambda: psser.astype(np.int32),
- )
- with ps.option_context("compute.eager_check", False):
- psser.astype(bool)
- psser.astype(int)
- psser.astype(np.int32)
-
- def test_neg(self):
- # pandas raises "TypeError: bad operand type for unary -:
'FloatingArray'"
- for dtype in self.fractional_extension_dtypes:
- self.assert_eq(
- ps.Series([-0.1, -0.2, -0.3, None], dtype=dtype),
- -ps.Series([0.1, 0.2, 0.3, None], dtype=dtype),
- )
-
- def test_abs(self):
- for pser, psser in self.fractional_extension_pser_psser_pairs:
- self.check_extension(abs(pser), abs(psser))
-
- def test_invert(self):
- for psser in self.fractional_extension_pssers:
- self.assertRaises(TypeError, lambda: ~psser)
-
- def test_eq(self):
- with option_context("compute.ops_on_diff_frames", True):
- for pser, psser in self.fractional_extension_pser_psser_pairs:
- self.check_extension(pser == pser, (psser ==
psser).sort_index())
-
- def test_ne(self):
- with option_context("compute.ops_on_diff_frames", True):
- for pser, psser in self.fractional_extension_pser_psser_pairs:
- self.check_extension(pser != pser, (psser !=
psser).sort_index())
-
- def test_lt(self):
- with option_context("compute.ops_on_diff_frames", True):
- for pser, psser in self.fractional_extension_pser_psser_pairs:
- self.check_extension(pser < pser, (psser < psser).sort_index())
-
- def test_le(self):
- with option_context("compute.ops_on_diff_frames", True):
- for pser, psser in self.fractional_extension_pser_psser_pairs:
- self.check_extension(pser <= pser, (psser <=
psser).sort_index())
-
- def test_gt(self):
- with option_context("compute.ops_on_diff_frames", True):
- for pser, psser in self.fractional_extension_pser_psser_pairs:
- self.check_extension(pser > pser, (psser > psser).sort_index())
-
- def test_ge(self):
- with option_context("compute.ops_on_diff_frames", True):
- for pser, psser in self.fractional_extension_pser_psser_pairs:
- self.check_extension(pser >= pser, (psser >=
psser).sort_index())
-
-
class NumOpsTests(
NumOpsTestsMixin,
OpsTestBase,
@@ -439,22 +201,6 @@ class NumOpsTests(
pass
-class IntegralExtensionOpsTests(
- IntegralExtensionOpsTestsMixin,
- OpsTestBase,
- PandasOnSparkTestCase,
-):
- pass
-
-
-class FractionalExtensionOpsTests(
- FractionalExtensionOpsTestsMixin,
- OpsTestBase,
- PandasOnSparkTestCase,
-):
- pass
-
-
if __name__ == "__main__":
from pyspark.testing import main
diff --git
a/python/pyspark/pandas/tests/data_type_ops/test_num_ops_fractional_ext.py
b/python/pyspark/pandas/tests/data_type_ops/test_num_ops_fractional_ext.py
new file mode 100644
index 000000000000..e507029ad896
--- /dev/null
+++ b/python/pyspark/pandas/tests/data_type_ops/test_num_ops_fractional_ext.py
@@ -0,0 +1,84 @@
+#
+# 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.testing.pandasutils import PandasOnSparkTestCase
+from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase
+from pyspark.pandas.typedef.typehints import extension_float_dtypes_available
+
+
[email protected](
+ not extension_float_dtypes_available, "pandas extension float dtypes are
not available"
+)
+class FractionalExtensionOpsTestsMixin:
+ @property
+ def fractional_extension_psers(self):
+ return [
+ pd.Series([0.1, 0.2, 0.3, None], dtype=dtype)
+ for dtype in self.fractional_extension_dtypes
+ ]
+
+ @property
+ def fractional_extension_pssers(self):
+ return [ps.from_pandas(pser) for pser in
self.fractional_extension_psers]
+
+ @property
+ def fractional_extension_pser_psser_pairs(self):
+ return zip(self.fractional_extension_psers,
self.fractional_extension_pssers)
+
+ def test_from_to_pandas(self):
+ for pser, psser in self.fractional_extension_pser_psser_pairs:
+ self.check_extension(pser, psser._to_pandas())
+ self.check_extension(ps.from_pandas(pser), psser)
+
+ def test_isnull(self):
+ for pser, psser in self.fractional_extension_pser_psser_pairs:
+ self.assert_eq(pser.isnull(), psser.isnull())
+
+ def test_neg(self):
+ # pandas raises "TypeError: bad operand type for unary -:
'FloatingArray'"
+ for dtype in self.fractional_extension_dtypes:
+ self.assert_eq(
+ ps.Series([-0.1, -0.2, -0.3, None], dtype=dtype),
+ -ps.Series([0.1, 0.2, 0.3, None], dtype=dtype),
+ )
+
+ def test_abs(self):
+ for pser, psser in self.fractional_extension_pser_psser_pairs:
+ self.check_extension(abs(pser), abs(psser))
+
+ def test_invert(self):
+ for psser in self.fractional_extension_pssers:
+ self.assertRaises(TypeError, lambda: ~psser)
+
+
+class FractionalExtensionOpsTests(
+ FractionalExtensionOpsTestsMixin,
+ OpsTestBase,
+ PandasOnSparkTestCase,
+):
+ pass
+
+
+if __name__ == "__main__":
+ from pyspark.testing import main
+
+ main()
diff --git
a/python/pyspark/pandas/tests/data_type_ops/test_num_ops_fractional_ext_astype_cmp.py
b/python/pyspark/pandas/tests/data_type_ops/test_num_ops_fractional_ext_astype_cmp.py
new file mode 100644
index 000000000000..aa429e6b4185
--- /dev/null
+++
b/python/pyspark/pandas/tests/data_type_ops/test_num_ops_fractional_ext_astype_cmp.py
@@ -0,0 +1,120 @@
+#
+# 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
+import numpy as np
+
+from pyspark import pandas as ps
+from pyspark.pandas.config import option_context
+from pyspark.testing.pandasutils import PandasOnSparkTestCase
+from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase
+from pyspark.pandas.typedef.typehints import extension_float_dtypes_available
+
+
[email protected](
+ not extension_float_dtypes_available, "pandas extension float dtypes are
not available"
+)
+class FractionalExtensionAstypeCmpOpsTestsMixin:
+ @property
+ def fractional_extension_psers(self):
+ return [
+ pd.Series([0.1, 0.2, 0.3, None], dtype=dtype)
+ for dtype in self.fractional_extension_dtypes
+ ]
+
+ @property
+ def fractional_extension_pssers(self):
+ return [ps.from_pandas(pser) for pser in
self.fractional_extension_psers]
+
+ @property
+ def fractional_extension_pser_psser_pairs(self):
+ return zip(self.fractional_extension_psers,
self.fractional_extension_pssers)
+
+ def test_astype(self):
+ for pser, psser in self.fractional_extension_pser_psser_pairs:
+ for dtype in self.extension_dtypes:
+ self.check_extension(pser.astype(dtype), psser.astype(dtype))
+ for pser, psser in self.fractional_extension_pser_psser_pairs:
+ self.assert_eq(pser.astype(float), psser.astype(float))
+ self.assert_eq(pser.astype("category"), psser.astype("category"))
+ self.assert_eq(pser.astype(np.float32), psser.astype(np.float32))
+ with ps.option_context("compute.eager_check", True):
+ self.assertRaisesRegex(
+ ValueError,
+ "Cannot convert fractions with missing values to bool",
+ lambda: psser.astype(bool),
+ )
+ self.assertRaisesRegex(
+ ValueError,
+ "Cannot convert fractions with missing values to integer",
+ lambda: psser.astype(int),
+ )
+ self.assertRaisesRegex(
+ ValueError,
+ "Cannot convert fractions with missing values to integer",
+ lambda: psser.astype(np.int32),
+ )
+ with ps.option_context("compute.eager_check", False):
+ psser.astype(bool)
+ psser.astype(int)
+ psser.astype(np.int32)
+
+ def test_eq(self):
+ with option_context("compute.ops_on_diff_frames", True):
+ for pser, psser in self.fractional_extension_pser_psser_pairs:
+ self.check_extension(pser == pser, (psser ==
psser).sort_index())
+
+ def test_ne(self):
+ with option_context("compute.ops_on_diff_frames", True):
+ for pser, psser in self.fractional_extension_pser_psser_pairs:
+ self.check_extension(pser != pser, (psser !=
psser).sort_index())
+
+ def test_lt(self):
+ with option_context("compute.ops_on_diff_frames", True):
+ for pser, psser in self.fractional_extension_pser_psser_pairs:
+ self.check_extension(pser < pser, (psser < psser).sort_index())
+
+ def test_le(self):
+ with option_context("compute.ops_on_diff_frames", True):
+ for pser, psser in self.fractional_extension_pser_psser_pairs:
+ self.check_extension(pser <= pser, (psser <=
psser).sort_index())
+
+ def test_gt(self):
+ with option_context("compute.ops_on_diff_frames", True):
+ for pser, psser in self.fractional_extension_pser_psser_pairs:
+ self.check_extension(pser > pser, (psser > psser).sort_index())
+
+ def test_ge(self):
+ with option_context("compute.ops_on_diff_frames", True):
+ for pser, psser in self.fractional_extension_pser_psser_pairs:
+ self.check_extension(pser >= pser, (psser >=
psser).sort_index())
+
+
+class FractionalExtensionAstypeCmpOpsTests(
+ FractionalExtensionAstypeCmpOpsTestsMixin,
+ OpsTestBase,
+ PandasOnSparkTestCase,
+):
+ pass
+
+
+if __name__ == "__main__":
+ from pyspark.testing import main
+
+ main()
diff --git
a/python/pyspark/pandas/tests/data_type_ops/test_num_ops_integral_ext.py
b/python/pyspark/pandas/tests/data_type_ops/test_num_ops_integral_ext.py
new file mode 100644
index 000000000000..6745b5a840f7
--- /dev/null
+++ b/python/pyspark/pandas/tests/data_type_ops/test_num_ops_integral_ext.py
@@ -0,0 +1,106 @@
+#
+# 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.testing.pandasutils import PandasOnSparkTestCase
+from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase
+from pyspark.pandas.typedef.typehints import extension_dtypes_available
+
+
[email protected](not extension_dtypes_available, "pandas extension dtypes are
not available")
+class IntegralExtensionOpsTestsMixin:
+ @property
+ def intergral_extension_psers(self):
+ return [pd.Series([1, 2, 3, None], dtype=dtype) for dtype in
self.integral_extension_dtypes]
+
+ @property
+ def intergral_extension_pssers(self):
+ return [ps.from_pandas(pser) for pser in
self.intergral_extension_psers]
+
+ @property
+ def intergral_extension_pser_psser_pairs(self):
+ return zip(self.intergral_extension_psers,
self.intergral_extension_pssers)
+
+ def test_from_to_pandas(self):
+ for pser, psser in self.intergral_extension_pser_psser_pairs:
+ self.check_extension(pser, psser._to_pandas())
+ self.check_extension(ps.from_pandas(pser), psser)
+
+ def test_isnull(self):
+ for pser, psser in self.intergral_extension_pser_psser_pairs:
+ self.assert_eq(pser.isnull(), psser.isnull())
+
+ def test_neg(self):
+ for pser, psser in self.intergral_extension_pser_psser_pairs:
+ self.check_extension(-pser, -psser)
+
+ def test_abs(self):
+ for pser, psser in self.intergral_extension_pser_psser_pairs:
+ self.check_extension(abs(pser), abs(psser))
+
+ def test_invert(self):
+ for pser, psser in self.intergral_extension_pser_psser_pairs:
+ self.check_extension(~pser, ~psser)
+
+ def test_xor(self):
+ for psser in self.intergral_extension_pssers:
+ self.assertRaisesRegex(
+ TypeError,
+ "XOR can not be applied to given types.",
+ lambda: psser ^ 1,
+ )
+ self.assertRaisesRegex(
+ TypeError,
+ "XOR can not be applied to given types.",
+ lambda: psser ^ psser,
+ )
+ self.assertRaisesRegex(
+ TypeError,
+ "XOR can not be applied to given types.",
+ lambda: psser ^ False,
+ )
+
+ def test_rxor(self):
+ for psser in self.intergral_extension_pssers:
+ self.assertRaisesRegex(
+ TypeError,
+ "XOR can not be applied to given types.",
+ lambda: 1 ^ psser,
+ )
+ self.assertRaisesRegex(
+ TypeError,
+ "XOR can not be applied to given types.",
+ lambda: False ^ psser,
+ )
+
+
+class IntegralExtensionOpsTests(
+ IntegralExtensionOpsTestsMixin,
+ OpsTestBase,
+ PandasOnSparkTestCase,
+):
+ pass
+
+
+if __name__ == "__main__":
+ from pyspark.testing import main
+
+ main()
diff --git
a/python/pyspark/pandas/tests/data_type_ops/test_num_ops_integral_ext_astype_cmp.py
b/python/pyspark/pandas/tests/data_type_ops/test_num_ops_integral_ext_astype_cmp.py
new file mode 100644
index 000000000000..9f8224eb0798
--- /dev/null
+++
b/python/pyspark/pandas/tests/data_type_ops/test_num_ops_integral_ext_astype_cmp.py
@@ -0,0 +1,117 @@
+#
+# 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
+import numpy as np
+
+from pyspark import pandas as ps
+from pyspark.pandas.config import option_context
+from pyspark.testing.pandasutils import PandasOnSparkTestCase
+from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase
+from pyspark.pandas.typedef.typehints import extension_dtypes_available
+
+
[email protected](not extension_dtypes_available, "pandas extension dtypes are
not available")
+class IntegralExtensionAstypeCmpOpsTestsMixin:
+ @property
+ def intergral_extension_psers(self):
+ return [pd.Series([1, 2, 3, None], dtype=dtype) for dtype in
self.integral_extension_dtypes]
+
+ @property
+ def intergral_extension_pssers(self):
+ return [ps.from_pandas(pser) for pser in
self.intergral_extension_psers]
+
+ @property
+ def intergral_extension_pser_psser_pairs(self):
+ return zip(self.intergral_extension_psers,
self.intergral_extension_pssers)
+
+ def test_astype(self):
+ for pser, psser in self.intergral_extension_pser_psser_pairs:
+ for dtype in self.extension_dtypes:
+ if dtype in self.string_extension_dtype:
+ self.check_extension(pser.astype(dtype),
psser.astype(dtype))
+ else:
+ self.check_extension(pser.astype(dtype),
psser.astype(dtype))
+ for pser, psser in self.intergral_extension_pser_psser_pairs:
+ self.assert_eq(pser.astype(float), psser.astype(float))
+ self.assert_eq(pser.astype(np.float32), psser.astype(np.float32))
+ with ps.option_context("compute.eager_check", True):
+ self.assertRaisesRegex(
+ ValueError,
+ "Cannot convert integrals with missing values to bool",
+ lambda: psser.astype(bool),
+ )
+ self.assertRaisesRegex(
+ ValueError,
+ "Cannot convert integrals with missing values to integer",
+ lambda: psser.astype(int),
+ )
+ self.assertRaisesRegex(
+ ValueError,
+ "Cannot convert integrals with missing values to integer",
+ lambda: psser.astype(np.int32),
+ )
+ with ps.option_context("compute.eager_check", False):
+ psser.astype(bool)
+ psser.astype(int)
+ psser.astype(np.int32)
+
+ def test_eq(self):
+ with option_context("compute.ops_on_diff_frames", True):
+ for pser, psser in self.intergral_extension_pser_psser_pairs:
+ self.check_extension(pser == pser, (psser ==
psser).sort_index())
+
+ def test_ne(self):
+ with option_context("compute.ops_on_diff_frames", True):
+ for pser, psser in self.intergral_extension_pser_psser_pairs:
+ self.check_extension(pser != pser, (psser !=
psser).sort_index())
+
+ def test_lt(self):
+ with option_context("compute.ops_on_diff_frames", True):
+ for pser, psser in self.intergral_extension_pser_psser_pairs:
+ self.check_extension(pser < pser, (psser < psser).sort_index())
+
+ def test_le(self):
+ with option_context("compute.ops_on_diff_frames", True):
+ for pser, psser in self.intergral_extension_pser_psser_pairs:
+ self.check_extension(pser <= pser, (psser <=
psser).sort_index())
+
+ def test_gt(self):
+ with option_context("compute.ops_on_diff_frames", True):
+ for pser, psser in self.intergral_extension_pser_psser_pairs:
+ self.check_extension(pser > pser, (psser > psser).sort_index())
+
+ def test_ge(self):
+ with option_context("compute.ops_on_diff_frames", True):
+ for pser, psser in self.intergral_extension_pser_psser_pairs:
+ self.check_extension(pser >= pser, (psser >=
psser).sort_index())
+
+
+class IntegralExtensionAstypeCmpOpsTests(
+ IntegralExtensionAstypeCmpOpsTestsMixin,
+ OpsTestBase,
+ PandasOnSparkTestCase,
+):
+ pass
+
+
+if __name__ == "__main__":
+ from pyspark.testing import main
+
+ main()
diff --git a/python/pyspark/pandas/tests/groupby/test_split_apply.py
b/python/pyspark/pandas/tests/groupby/test_split_apply.py
index 27835788e762..d23ff2b49a71 100644
--- a/python/pyspark/pandas/tests/groupby/test_split_apply.py
+++ b/python/pyspark/pandas/tests/groupby/test_split_apply.py
@@ -160,7 +160,6 @@ class
GroupbySplitApplyMixin(GroupbySplitApplyTestingFuncMixin):
def test_split_apply_combine_on_series(self):
funcs = [
((True, False), ["sum"]),
- ((True, True), ["mean"]),
]
self._test_split_apply_func(funcs)
diff --git
a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
b/python/pyspark/pandas/tests/groupby/test_split_apply_mean.py
similarity index 52%
copy from
python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
copy to python/pyspark/pandas/tests/groupby/test_split_apply_mean.py
index ed8dd1083c48..dd260c5dd187 100644
--- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
+++ b/python/pyspark/pandas/tests/groupby/test_split_apply_mean.py
@@ -15,39 +15,21 @@
# limitations under the License.
#
-from pyspark.pandas.tests.data_type_ops.test_num_ops import (
- NumOpsTestsMixin,
- IntegralExtensionOpsTestsMixin,
- FractionalExtensionOpsTestsMixin,
-)
-from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase
-from pyspark.testing.pandasutils import PandasOnSparkTestUtils
-from pyspark.testing.connectutils import ReusedConnectTestCase
-
-
-class NumOpsParityTests(
- NumOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
- ReusedConnectTestCase,
-):
- pass
+from pyspark.testing.pandasutils import PandasOnSparkTestCase
+from pyspark.pandas.tests.groupby.test_split_apply import
GroupbySplitApplyTestingFuncMixin
-class IntegralExtensionOpsParityTests(
- IntegralExtensionOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
- ReusedConnectTestCase,
-):
- pass
+class GroupbySplitApplyMeanMixin(GroupbySplitApplyTestingFuncMixin):
+ def test_split_apply_combine_on_series(self):
+ funcs = [
+ ((True, True), ["mean"]),
+ ]
+ self._test_split_apply_func(funcs)
-class FractionalExtensionOpsParityTests(
- FractionalExtensionOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
- ReusedConnectTestCase,
+class GroupbySplitApplyMeanTests(
+ GroupbySplitApplyMeanMixin,
+ PandasOnSparkTestCase,
):
pass
diff --git a/python/pyspark/pandas/tests/groupby/test_stat.py
b/python/pyspark/pandas/tests/groupby/test_stat.py
index d1f44131db71..ae715407ecf9 100644
--- a/python/pyspark/pandas/tests/groupby/test_stat.py
+++ b/python/pyspark/pandas/tests/groupby/test_stat.py
@@ -124,37 +124,6 @@ class GroupbyStatMixin(GroupbyStatTestingFuncMixin):
expected_error=ValueError if using_pandas3 else None,
)
- def test_median(self):
- psdf = ps.DataFrame(
- {
- "a": [1.0, 1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0],
- "b": [2.0, 3.0, 1.0, 4.0, 6.0, 9.0, 8.0, 10.0, 7.0, 5.0],
- "c": [3.0, 5.0, 2.0, 5.0, 1.0, 2.0, 6.0, 4.0, 3.0, 6.0],
- },
- columns=["a", "b", "c"],
- index=[7, 2, 4, 1, 3, 4, 9, 10, 5, 6],
- )
- # DataFrame
- expected_result = ps.DataFrame(
- {"b": [2.0, 8.0, 7.0], "c": [3.0, 2.0, 4.0]}, index=pd.Index([1.0,
2.0, 3.0], name="a")
- )
- self.assert_eq(expected_result,
psdf.groupby("a").median().sort_index())
- # Series
- expected_result = ps.Series(
- [2.0, 8.0, 7.0], name="b", index=pd.Index([1.0, 2.0, 3.0],
name="a")
- )
- self.assert_eq(expected_result,
psdf.groupby("a")["b"].median().sort_index())
-
- with self.assertRaisesRegex(TypeError, "accuracy must be an integer;
however"):
- psdf.groupby("a").median(accuracy="a")
-
- if LooseVersion(pd.__version__) >= "3.0.0":
- # pandas < 3 raises an error when numeric_only is False or None
- self._test_stat_func(
- lambda groupby_obj: groupby_obj.median(numeric_only=None),
- expected_error=ValueError if using_pandas3 else None,
- )
-
class GroupbyStatTests(
GroupbyStatMixin,
diff --git a/python/pyspark/pandas/tests/groupby/test_stat_median.py
b/python/pyspark/pandas/tests/groupby/test_stat_median.py
new file mode 100644
index 000000000000..a555e4b06e44
--- /dev/null
+++ b/python/pyspark/pandas/tests/groupby/test_stat_median.py
@@ -0,0 +1,84 @@
+#
+# 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.loose_version import LooseVersion
+from pyspark.testing.pandasutils import PandasOnSparkTestCase
+from pyspark.pandas.tests.groupby.test_stat import
GroupbyStatTestingFuncMixin, using_pandas3
+
+
+class GroupbyStatMedianMixin(GroupbyStatTestingFuncMixin):
+ @property
+ def pdf(self):
+ return pd.DataFrame(
+ {
+ "A": [1, 2, 1, 2],
+ "B": [3.1, 4.1, 4.1, 3.1],
+ "C": ["a", "b", "b", "a"],
+ "D": [True, False, False, True],
+ }
+ )
+
+ @property
+ def psdf(self):
+ return ps.from_pandas(self.pdf)
+
+ def test_median(self):
+ psdf = ps.DataFrame(
+ {
+ "a": [1.0, 1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0],
+ "b": [2.0, 3.0, 1.0, 4.0, 6.0, 9.0, 8.0, 10.0, 7.0, 5.0],
+ "c": [3.0, 5.0, 2.0, 5.0, 1.0, 2.0, 6.0, 4.0, 3.0, 6.0],
+ },
+ columns=["a", "b", "c"],
+ index=[7, 2, 4, 1, 3, 4, 9, 10, 5, 6],
+ )
+ # DataFrame
+ expected_result = ps.DataFrame(
+ {"b": [2.0, 8.0, 7.0], "c": [3.0, 2.0, 4.0]}, index=pd.Index([1.0,
2.0, 3.0], name="a")
+ )
+ self.assert_eq(expected_result,
psdf.groupby("a").median().sort_index())
+ # Series
+ expected_result = ps.Series(
+ [2.0, 8.0, 7.0], name="b", index=pd.Index([1.0, 2.0, 3.0],
name="a")
+ )
+ self.assert_eq(expected_result,
psdf.groupby("a")["b"].median().sort_index())
+
+ with self.assertRaisesRegex(TypeError, "accuracy must be an integer;
however"):
+ psdf.groupby("a").median(accuracy="a")
+
+ if LooseVersion(pd.__version__) >= "3.0.0":
+ # pandas < 3 raises an error when numeric_only is False or None
+ self._test_stat_func(
+ lambda groupby_obj: groupby_obj.median(numeric_only=None),
+ expected_error=ValueError if using_pandas3 else None,
+ )
+
+
+class GroupbyStatMedianTests(
+ GroupbyStatMedianMixin,
+ PandasOnSparkTestCase,
+):
+ pass
+
+
+if __name__ == "__main__":
+ from pyspark.testing import main
+
+ main()
diff --git a/python/pyspark/sql/tests/arrow/test_arrow_cogrouped_map.py
b/python/pyspark/sql/tests/arrow/test_arrow_cogrouped_map.py
index 5b272f89bb5d..5db694fcd953 100644
--- a/python/pyspark/sql/tests/arrow/test_arrow_cogrouped_map.py
+++ b/python/pyspark/sql/tests/arrow/test_arrow_cogrouped_map.py
@@ -17,14 +17,12 @@
import os
import time
import unittest
-import logging
from pyspark.errors import PythonException
from pyspark.sql import Row
from pyspark.sql import functions as sf
from pyspark.testing.sqlutils import ReusedSQLTestCase
-from pyspark.testing.utils import assertDataFrameEqual, have_pyarrow,
pyarrow_requirement_message
-from pyspark.util import is_remote_only
+from pyspark.testing.utils import have_pyarrow, pyarrow_requirement_message
if have_pyarrow:
import pyarrow as pa
@@ -35,7 +33,7 @@ if have_pyarrow:
not have_pyarrow,
pyarrow_requirement_message,
)
-class CogroupedMapInArrowTestsMixin:
+class CogroupedMapInArrowTestsFuncMixin:
@property
def left(self):
return self.spark.range(0, 10, 2, 3).withColumn("v", sf.col("id") * 10)
@@ -77,14 +75,14 @@ class CogroupedMapInArrowTestsMixin:
for table in [left, right]
for k in table.column(key_column)
)
- return CogroupedMapInArrowTestsMixin.apply_in_arrow_func(left,
right)
+ return CogroupedMapInArrowTestsFuncMixin.apply_in_arrow_func(left,
right)
return func
@staticmethod
def apply_in_pandas_with_key_func(key_column):
def func(key, left, right):
- return
CogroupedMapInArrowTestsMixin.apply_in_arrow_with_key_func(key_column)(
+ return
CogroupedMapInArrowTestsFuncMixin.apply_in_arrow_with_key_func(key_column)(
tuple(pa.scalar(k) for k in key),
pa.Table.from_pandas(left),
pa.Table.from_pandas(right),
@@ -97,21 +95,23 @@ class CogroupedMapInArrowTestsMixin:
# compare with result of applyInPandas
expected = cogrouped_df.applyInPandas(
-
CogroupedMapInArrowTestsMixin.apply_in_pandas_with_key_func(key_column), schema
+
CogroupedMapInArrowTestsFuncMixin.apply_in_pandas_with_key_func(key_column),
schema
)
# apply in arrow without key
actual = cogrouped_df.applyInArrow(
- CogroupedMapInArrowTestsMixin.apply_in_arrow_func, schema
+ CogroupedMapInArrowTestsFuncMixin.apply_in_arrow_func, schema
).collect()
self.assertEqual(actual, expected.collect())
# apply in arrow with key
actual2 = cogrouped_df.applyInArrow(
-
CogroupedMapInArrowTestsMixin.apply_in_arrow_with_key_func(key_column), schema
+
CogroupedMapInArrowTestsFuncMixin.apply_in_arrow_with_key_func(key_column),
schema
).collect()
self.assertEqual(actual2, expected.collect())
+
+class CogroupedMapInArrowTestsMixin(CogroupedMapInArrowTestsFuncMixin):
def test_apply_in_arrow(self):
self.do_test_apply_in_arrow(self.cogrouped)
@@ -121,37 +121,6 @@ class CogroupedMapInArrowTestsMixin:
cogrouped_df = grouped_left_df.cogroup(grouped_right_df)
self.do_test_apply_in_arrow(cogrouped_df, key_column=None)
- def test_apply_in_arrow_large_var_types(self):
- # SPARK-56929: when useLargeVarTypes=true, the expected schema
computed by
- # worker.py for result validation must also use
large_string/large_binary,
- # otherwise verify_arrow_result raises a spurious
RESULT_COLUMN_TYPES_MISMATCH.
- left = self.spark.createDataFrame(
- [(0, "foo", b"foo"), (1, None, None)], "id long, s string, b
binary"
- )
- right = self.spark.createDataFrame(
- [(0, "bar", b"bar"), (1, "baz", b"baz")], "id long, s string, b
binary"
- )
- schema = "s string, b binary"
-
- def func(left_tbl, right_tbl):
- assert pa.types.is_large_string(left_tbl.schema.field("s").type)
- assert pa.types.is_large_binary(left_tbl.schema.field("b").type)
- return left_tbl.select(["s", "b"])
-
- expected = left.select("s", "b")
- for assign_cols_by_name in [True, False]:
- with self.subTest(assign_cols_by_name=assign_cols_by_name):
- with self.sql_conf(
- {
- "spark.sql.execution.arrow.useLargeVarTypes": True,
- "spark.sql.legacy.execution.pandas.groupedMap."
- "assignColumnsByName": assign_cols_by_name,
- }
- ):
- cogrouped = left.groupBy("id").cogroup(right.groupBy("id"))
- actual = cogrouped.applyInArrow(func, schema)
- assertDataFrameEqual(actual, expected)
-
def test_apply_in_arrow_not_returning_arrow_table(self):
def func(key, left, right):
return key
@@ -323,150 +292,6 @@ class CogroupedMapInArrowTestsMixin:
self.assertEqual(r.a, "hi")
self.assertEqual(r.b, 1)
- def test_with_local_data(self):
- df1 = self.spark.createDataFrame(
- [(1, 1.0, "a"), (2, 2.0, "b"), (1, 3.0, "c"), (2, 4.0, "d")],
("id", "v1", "v2")
- )
- df2 = self.spark.createDataFrame([(1, "x"), (2, "y"), (1, "z")],
("id", "v3"))
-
- def summarize(left, right):
- return pa.Table.from_pydict(
- {
- "left_rows": [left.num_rows],
- "left_columns": [left.num_columns],
- "right_rows": [right.num_rows],
- "right_columns": [right.num_columns],
- }
- )
-
- df = (
- df1.groupby("id")
- .cogroup(df2.groupby("id"))
- .applyInArrow(
- summarize,
- schema="left_rows long, left_columns long, right_rows long,
right_columns long",
- )
- )
-
- self.assertEqual(
- df._show_string(),
- "+---------+------------+----------+-------------+\n"
- "|left_rows|left_columns|right_rows|right_columns|\n"
- "+---------+------------+----------+-------------+\n"
- "| 2| 3| 2| 2|\n"
- "| 2| 3| 1| 2|\n"
- "+---------+------------+----------+-------------+\n",
- )
-
- def test_self_join(self):
- df = self.spark.createDataFrame([(1, 1)], ("k", "v"))
-
- def arrow_func(key, left, right):
- return pa.Table.from_pydict({"x": [2], "y": [2]})
-
- df2 =
df.groupby("k").cogroup(df.groupby("k")).applyInArrow(arrow_func, "x long, y
long")
-
- self.assertEqual(df2.join(df2).count(), 1)
-
- def test_arrow_batch_slicing(self):
- m, n = 100000, 10000
-
- df1 = self.spark.range(m).select((sf.col("id") % 2).alias("key"),
sf.col("id").alias("v"))
- cols = {f"col_{i}": sf.col("v") + i for i in range(10)}
- df1 = df1.withColumns(cols)
-
- df2 = self.spark.range(n).select((sf.col("id") % 4).alias("key"),
sf.col("id").alias("v"))
- cols = {f"col_{i}": sf.col("v") + i for i in range(20)}
- df2 = df2.withColumns(cols)
-
- def summarize(key, left, right):
- assert len(left) == m / 2 or len(left) == 0, len(left)
- assert len(right) == n / 4, len(right)
- return pa.Table.from_pydict(
- {
- "key": [key[0].as_py()],
- "left_rows": [left.num_rows],
- "left_columns": [left.num_columns],
- "right_rows": [right.num_rows],
- "right_columns": [right.num_columns],
- }
- )
-
- schema = "key long, left_rows long, left_columns long, right_rows
long, right_columns long"
-
- expected = [
- Row(key=0, left_rows=m / 2, left_columns=12, right_rows=n / 4,
right_columns=22),
- Row(key=1, left_rows=m / 2, left_columns=12, right_rows=n / 4,
right_columns=22),
- Row(key=2, left_rows=0, left_columns=12, right_rows=n / 4,
right_columns=22),
- Row(key=3, left_rows=0, left_columns=12, right_rows=n / 4,
right_columns=22),
- ]
-
- for maxRecords, maxBytes in [(1000, 2**31 - 1), (0, 4096), (1000,
4096)]:
- with self.subTest(maxRecords=maxRecords, maxBytes=maxBytes):
- with self.sql_conf(
- {
- "spark.sql.execution.arrow.maxRecordsPerBatch":
maxRecords,
- "spark.sql.execution.arrow.maxBytesPerBatch": maxBytes,
- }
- ):
- result = (
- df1.groupby("key")
- .cogroup(df2.groupby("key"))
- .applyInArrow(summarize, schema=schema)
- .sort("key")
- .collect()
- )
-
- self.assertEqual(expected, result)
-
- def test_negative_and_zero_batch_size(self):
- for batch_size in [0, -1]:
- with
self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": batch_size}):
- CogroupedMapInArrowTestsMixin.test_apply_in_arrow(self)
-
- @unittest.skipIf(is_remote_only(), "Requires JVM access")
- def test_cogroup_apply_in_arrow_with_logging(self):
- import pyarrow as pa
-
- def func_with_logging(left, right):
- assert isinstance(left, pa.Table)
- assert isinstance(right, pa.Table)
- logger = logging.getLogger("test_arrow_cogrouped_map")
- logger.warning(
- "arrow cogrouped map: "
- + f"{dict(v1=left['v1'].to_pylist(),
v2=right['v2'].to_pylist())}"
- )
- return left.join(right, keys="id", join_type="inner")
-
- left_df = self.spark.createDataFrame([(1, 10), (2, 20), (1, 30)],
["id", "v1"])
- right_df = self.spark.createDataFrame([(1, 100), (2, 200), (1, 300)],
["id", "v2"])
-
- grouped_left = left_df.groupBy("id")
- grouped_right = right_df.groupBy("id")
- cogrouped_df = grouped_left.cogroup(grouped_right)
-
- with self.sql_conf({"spark.sql.pyspark.worker.logging.enabled":
"true"}):
- assertDataFrameEqual(
- cogrouped_df.applyInArrow(func_with_logging, "id long, v1
long, v2 long"),
- [Row(id=1, v1=v1, v2=v2) for v1 in [10, 30] for v2 in [100,
300]]
- + [Row(id=2, v1=20, v2=200)],
- )
-
- logs = self.spark.tvf.python_worker_logs()
-
- assertDataFrameEqual(
- logs.select("level", "msg", "context", "logger"),
- [
- Row(
- level="WARNING",
- msg=f"arrow cogrouped map: {dict(v1=v1, v2=v2)}",
- context={"func_name": func_with_logging.__name__},
- logger="test_arrow_cogrouped_map",
- )
- for v1, v2 in [([10, 30], [100, 300]), ([20], [200])]
- ],
- )
-
class CogroupedMapInArrowTests(CogroupedMapInArrowTestsMixin,
ReusedSQLTestCase):
@classmethod
diff --git a/python/pyspark/sql/tests/arrow/test_arrow_cogrouped_map_misc.py
b/python/pyspark/sql/tests/arrow/test_arrow_cogrouped_map_misc.py
new file mode 100644
index 000000000000..15eb0b2320cb
--- /dev/null
+++ b/python/pyspark/sql/tests/arrow/test_arrow_cogrouped_map_misc.py
@@ -0,0 +1,240 @@
+#
+# 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 os
+import time
+import unittest
+import logging
+
+from pyspark.sql import Row
+from pyspark.sql import functions as sf
+from pyspark.sql.tests.arrow.test_arrow_cogrouped_map import
CogroupedMapInArrowTestsFuncMixin
+from pyspark.testing.sqlutils import ReusedSQLTestCase
+from pyspark.testing.utils import assertDataFrameEqual, have_pyarrow,
pyarrow_requirement_message
+from pyspark.util import is_remote_only
+
+if have_pyarrow:
+ import pyarrow as pa
+
+
[email protected](
+ not have_pyarrow,
+ pyarrow_requirement_message,
+)
+class CogroupedMapInArrowMiscTestsMixin(CogroupedMapInArrowTestsFuncMixin):
+ def test_apply_in_arrow_large_var_types(self):
+ # SPARK-56929: when useLargeVarTypes=true, the expected schema
computed by
+ # worker.py for result validation must also use
large_string/large_binary,
+ # otherwise verify_arrow_result raises a spurious
RESULT_COLUMN_TYPES_MISMATCH.
+ left = self.spark.createDataFrame(
+ [(0, "foo", b"foo"), (1, None, None)], "id long, s string, b
binary"
+ )
+ right = self.spark.createDataFrame(
+ [(0, "bar", b"bar"), (1, "baz", b"baz")], "id long, s string, b
binary"
+ )
+ schema = "s string, b binary"
+
+ def func(left_tbl, right_tbl):
+ assert pa.types.is_large_string(left_tbl.schema.field("s").type)
+ assert pa.types.is_large_binary(left_tbl.schema.field("b").type)
+ return left_tbl.select(["s", "b"])
+
+ expected = left.select("s", "b")
+ for assign_cols_by_name in [True, False]:
+ with self.subTest(assign_cols_by_name=assign_cols_by_name):
+ with self.sql_conf(
+ {
+ "spark.sql.execution.arrow.useLargeVarTypes": True,
+ "spark.sql.legacy.execution.pandas.groupedMap."
+ "assignColumnsByName": assign_cols_by_name,
+ }
+ ):
+ cogrouped = left.groupBy("id").cogroup(right.groupBy("id"))
+ actual = cogrouped.applyInArrow(func, schema)
+ assertDataFrameEqual(actual, expected)
+
+ def test_with_local_data(self):
+ df1 = self.spark.createDataFrame(
+ [(1, 1.0, "a"), (2, 2.0, "b"), (1, 3.0, "c"), (2, 4.0, "d")],
("id", "v1", "v2")
+ )
+ df2 = self.spark.createDataFrame([(1, "x"), (2, "y"), (1, "z")],
("id", "v3"))
+
+ def summarize(left, right):
+ return pa.Table.from_pydict(
+ {
+ "left_rows": [left.num_rows],
+ "left_columns": [left.num_columns],
+ "right_rows": [right.num_rows],
+ "right_columns": [right.num_columns],
+ }
+ )
+
+ df = (
+ df1.groupby("id")
+ .cogroup(df2.groupby("id"))
+ .applyInArrow(
+ summarize,
+ schema="left_rows long, left_columns long, right_rows long,
right_columns long",
+ )
+ )
+
+ self.assertEqual(
+ df._show_string(),
+ "+---------+------------+----------+-------------+\n"
+ "|left_rows|left_columns|right_rows|right_columns|\n"
+ "+---------+------------+----------+-------------+\n"
+ "| 2| 3| 2| 2|\n"
+ "| 2| 3| 1| 2|\n"
+ "+---------+------------+----------+-------------+\n",
+ )
+
+ def test_self_join(self):
+ df = self.spark.createDataFrame([(1, 1)], ("k", "v"))
+
+ def arrow_func(key, left, right):
+ return pa.Table.from_pydict({"x": [2], "y": [2]})
+
+ df2 =
df.groupby("k").cogroup(df.groupby("k")).applyInArrow(arrow_func, "x long, y
long")
+
+ self.assertEqual(df2.join(df2).count(), 1)
+
+ def test_arrow_batch_slicing(self):
+ m, n = 100000, 10000
+
+ df1 = self.spark.range(m).select((sf.col("id") % 2).alias("key"),
sf.col("id").alias("v"))
+ cols = {f"col_{i}": sf.col("v") + i for i in range(10)}
+ df1 = df1.withColumns(cols)
+
+ df2 = self.spark.range(n).select((sf.col("id") % 4).alias("key"),
sf.col("id").alias("v"))
+ cols = {f"col_{i}": sf.col("v") + i for i in range(20)}
+ df2 = df2.withColumns(cols)
+
+ def summarize(key, left, right):
+ assert len(left) == m / 2 or len(left) == 0, len(left)
+ assert len(right) == n / 4, len(right)
+ return pa.Table.from_pydict(
+ {
+ "key": [key[0].as_py()],
+ "left_rows": [left.num_rows],
+ "left_columns": [left.num_columns],
+ "right_rows": [right.num_rows],
+ "right_columns": [right.num_columns],
+ }
+ )
+
+ schema = "key long, left_rows long, left_columns long, right_rows
long, right_columns long"
+
+ expected = [
+ Row(key=0, left_rows=m / 2, left_columns=12, right_rows=n / 4,
right_columns=22),
+ Row(key=1, left_rows=m / 2, left_columns=12, right_rows=n / 4,
right_columns=22),
+ Row(key=2, left_rows=0, left_columns=12, right_rows=n / 4,
right_columns=22),
+ Row(key=3, left_rows=0, left_columns=12, right_rows=n / 4,
right_columns=22),
+ ]
+
+ for maxRecords, maxBytes in [(1000, 2**31 - 1), (0, 4096), (1000,
4096)]:
+ with self.subTest(maxRecords=maxRecords, maxBytes=maxBytes):
+ with self.sql_conf(
+ {
+ "spark.sql.execution.arrow.maxRecordsPerBatch":
maxRecords,
+ "spark.sql.execution.arrow.maxBytesPerBatch": maxBytes,
+ }
+ ):
+ result = (
+ df1.groupby("key")
+ .cogroup(df2.groupby("key"))
+ .applyInArrow(summarize, schema=schema)
+ .sort("key")
+ .collect()
+ )
+
+ self.assertEqual(expected, result)
+
+ def test_negative_and_zero_batch_size(self):
+ for batch_size in [0, -1]:
+ with
self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": batch_size}):
+ self.do_test_apply_in_arrow(self.cogrouped)
+
+ @unittest.skipIf(is_remote_only(), "Requires JVM access")
+ def test_cogroup_apply_in_arrow_with_logging(self):
+ import pyarrow as pa
+
+ def func_with_logging(left, right):
+ assert isinstance(left, pa.Table)
+ assert isinstance(right, pa.Table)
+ logger = logging.getLogger("test_arrow_cogrouped_map")
+ logger.warning(
+ "arrow cogrouped map: "
+ + f"{dict(v1=left['v1'].to_pylist(),
v2=right['v2'].to_pylist())}"
+ )
+ return left.join(right, keys="id", join_type="inner")
+
+ left_df = self.spark.createDataFrame([(1, 10), (2, 20), (1, 30)],
["id", "v1"])
+ right_df = self.spark.createDataFrame([(1, 100), (2, 200), (1, 300)],
["id", "v2"])
+
+ grouped_left = left_df.groupBy("id")
+ grouped_right = right_df.groupBy("id")
+ cogrouped_df = grouped_left.cogroup(grouped_right)
+
+ with self.sql_conf({"spark.sql.pyspark.worker.logging.enabled":
"true"}):
+ assertDataFrameEqual(
+ cogrouped_df.applyInArrow(func_with_logging, "id long, v1
long, v2 long"),
+ [Row(id=1, v1=v1, v2=v2) for v1 in [10, 30] for v2 in [100,
300]]
+ + [Row(id=2, v1=20, v2=200)],
+ )
+
+ logs = self.spark.tvf.python_worker_logs()
+
+ assertDataFrameEqual(
+ logs.select("level", "msg", "context", "logger"),
+ [
+ Row(
+ level="WARNING",
+ msg=f"arrow cogrouped map: {dict(v1=v1, v2=v2)}",
+ context={"func_name": func_with_logging.__name__},
+ logger="test_arrow_cogrouped_map",
+ )
+ for v1, v2 in [([10, 30], [100, 300]), ([20], [200])]
+ ],
+ )
+
+
+class CogroupedMapInArrowMiscTests(CogroupedMapInArrowMiscTestsMixin,
ReusedSQLTestCase):
+ @classmethod
+ def setUpClass(cls):
+ ReusedSQLTestCase.setUpClass()
+
+ # Synchronize default timezone between Python and Java
+ cls.tz_prev = os.environ.get("TZ", None) # save current tz if set
+ tz = "America/Los_Angeles"
+ os.environ["TZ"] = tz
+ time.tzset()
+
+ cls.sc.environment["TZ"] = tz
+ cls.spark.conf.set("spark.sql.session.timeZone", tz)
+
+ @classmethod
+ def tearDownClass(cls):
+ del os.environ["TZ"]
+ if cls.tz_prev is not None:
+ os.environ["TZ"] = cls.tz_prev
+ time.tzset()
+ ReusedSQLTestCase.tearDownClass()
+
+
+if __name__ == "__main__":
+ from pyspark.testing import main
+
+ main()
diff --git
a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
b/python/pyspark/sql/tests/connect/arrow/test_parity_arrow_cogrouped_map_misc.py
similarity index 55%
copy from
python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
copy to
python/pyspark/sql/tests/connect/arrow/test_parity_arrow_cogrouped_map_misc.py
index ed8dd1083c48..2a1922e55db5 100644
--- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
+++
b/python/pyspark/sql/tests/connect/arrow/test_parity_arrow_cogrouped_map_misc.py
@@ -15,40 +15,12 @@
# limitations under the License.
#
-from pyspark.pandas.tests.data_type_ops.test_num_ops import (
- NumOpsTestsMixin,
- IntegralExtensionOpsTestsMixin,
- FractionalExtensionOpsTestsMixin,
-)
-from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase
-from pyspark.testing.pandasutils import PandasOnSparkTestUtils
-from pyspark.testing.connectutils import ReusedConnectTestCase
-
-
-class NumOpsParityTests(
- NumOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
- ReusedConnectTestCase,
-):
- pass
-
-class IntegralExtensionOpsParityTests(
- IntegralExtensionOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
- ReusedConnectTestCase,
-):
- pass
+from pyspark.sql.tests.arrow.test_arrow_cogrouped_map_misc import
CogroupedMapInArrowMiscTestsMixin
+from pyspark.testing.connectutils import ReusedConnectTestCase
-class FractionalExtensionOpsParityTests(
- FractionalExtensionOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
- ReusedConnectTestCase,
-):
+class CogroupedMapInArrowMiscParityTests(CogroupedMapInArrowMiscTestsMixin,
ReusedConnectTestCase):
pass
diff --git
a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
b/python/pyspark/sql/tests/connect/pandas/test_parity_pandas_cogrouped_map_misc.py
similarity index 57%
copy from
python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
copy to
python/pyspark/sql/tests/connect/pandas/test_parity_pandas_cogrouped_map_misc.py
index ed8dd1083c48..532accd2707f 100644
--- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py
+++
b/python/pyspark/sql/tests/connect/pandas/test_parity_pandas_cogrouped_map_misc.py
@@ -15,38 +15,14 @@
# limitations under the License.
#
-from pyspark.pandas.tests.data_type_ops.test_num_ops import (
- NumOpsTestsMixin,
- IntegralExtensionOpsTestsMixin,
- FractionalExtensionOpsTestsMixin,
+from pyspark.sql.tests.pandas.test_pandas_cogrouped_map_misc import (
+ CogroupedApplyInPandasMiscTestsMixin,
)
-from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase
-from pyspark.testing.pandasutils import PandasOnSparkTestUtils
from pyspark.testing.connectutils import ReusedConnectTestCase
-class NumOpsParityTests(
- NumOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
- ReusedConnectTestCase,
-):
- pass
-
-
-class IntegralExtensionOpsParityTests(
- IntegralExtensionOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
- ReusedConnectTestCase,
-):
- pass
-
-
-class FractionalExtensionOpsParityTests(
- FractionalExtensionOpsTestsMixin,
- PandasOnSparkTestUtils,
- OpsTestBase,
+class CogroupedApplyInPandasMiscTests(
+ CogroupedApplyInPandasMiscTestsMixin,
ReusedConnectTestCase,
):
pass
diff --git a/python/pyspark/sql/tests/pandas/test_pandas_cogrouped_map.py
b/python/pyspark/sql/tests/pandas/test_pandas_cogrouped_map.py
index 625341ad0b51..b5a4daf3279c 100644
--- a/python/pyspark/sql/tests/pandas/test_pandas_cogrouped_map.py
+++ b/python/pyspark/sql/tests/pandas/test_pandas_cogrouped_map.py
@@ -16,7 +16,6 @@
#
import unittest
-import logging
from pyspark.loose_version import LooseVersion
from pyspark.sql import functions as sf
@@ -30,17 +29,14 @@ from pyspark.sql.types import (
YearMonthIntervalType,
Row,
)
-from pyspark.sql.window import Window
from pyspark.errors import IllegalArgumentException, PythonException
from pyspark.testing.sqlutils import ReusedSQLTestCase
from pyspark.testing.utils import (
- assertDataFrameEqual,
have_pandas,
have_pyarrow,
pandas_requirement_message,
pyarrow_requirement_message,
)
-from pyspark.util import is_remote_only
if have_pandas:
import pandas as pd
@@ -446,98 +442,6 @@ class CogroupedApplyInPandasTestsMixin:
self.assertEqual(row.asDict(), Row(column=2, value=2).asDict())
- def test_with_window_function(self):
- # SPARK-42168: a window function with same partition keys but
differing key order
- ids = 2
- days = 100
- vals = 10000
- parts = 10
-
- id_df = self.spark.range(ids)
- day_df = self.spark.range(days).withColumnRenamed("id", "day")
- vals_df = self.spark.range(vals).withColumnRenamed("id", "value")
- df = id_df.join(day_df).join(vals_df)
-
- left_df = df.withColumnRenamed("value",
"left").repartition(parts).cache()
- # SPARK-42132: this bug requires us to alias all columns from df here
- right_df = (
- df.select(
- sf.col("id").alias("id"), sf.col("day").alias("day"),
sf.col("value").alias("right")
- )
- .repartition(parts)
- .cache()
- )
-
- # note the column order is different to the groupBy("id", "day")
column order below
- window = Window.partitionBy("day", "id")
-
- left_grouped_df = left_df.groupBy("id", "day")
- right_grouped_df = right_df.withColumn(
- "day_sum", sf.sum(sf.col("day")).over(window)
- ).groupBy("id", "day")
-
- def cogroup(left: pd.DataFrame, right: pd.DataFrame) -> pd.DataFrame:
- return pd.DataFrame(
- [
- {
- "id": (
- left["id"][0]
- if not left.empty
- else (right["id"][0] if not right.empty else None)
- ),
- "day": (
- left["day"][0]
- if not left.empty
- else (right["day"][0] if not right.empty else None)
- ),
- "lefts": len(left.index),
- "rights": len(right.index),
- }
- ]
- )
-
- df = left_grouped_df.cogroup(right_grouped_df).applyInPandas(
- cogroup, schema="id long, day long, lefts integer, rights integer"
- )
-
- actual = df.orderBy("id", "day").take(days)
- self.assertEqual(actual, [Row(0, day, vals, vals) for day in
range(days)])
-
- def test_with_local_data(self):
- df1 = self.spark.createDataFrame(
- [(1, 1.0, "a"), (2, 2.0, "b"), (1, 3.0, "c"), (2, 4.0, "d")],
("id", "v1", "v2")
- )
- df2 = self.spark.createDataFrame([(1, "x"), (2, "y"), (1, "z")],
("id", "v3"))
-
- def summarize(left, right):
- return pd.DataFrame(
- {
- "left_rows": [len(left)],
- "left_columns": [len(left.columns)],
- "right_rows": [len(right)],
- "right_columns": [len(right.columns)],
- }
- )
-
- df = (
- df1.groupby("id")
- .cogroup(df2.groupby("id"))
- .applyInPandas(
- summarize,
- schema="left_rows long, left_columns long, right_rows long,
right_columns long",
- )
- )
-
- self.assertEqual(
- df._show_string(),
- "+---------+------------+----------+-------------+\n"
- "|left_rows|left_columns|right_rows|right_columns|\n"
- "+---------+------------+----------+-------------+\n"
- "| 2| 3| 2| 2|\n"
- "| 2| 3| 1| 2|\n"
- "+---------+------------+----------+-------------+\n",
- )
-
@staticmethod
def _test_with_key(left, right, isLeft):
def right_assign_key(key, lft, rgt):
@@ -664,104 +568,11 @@ class CogroupedApplyInPandasTestsMixin:
with self.assertRaisesRegex(errorClass, error_message_regex):
self.__test_merge(left, right, by, fn, output_schema)
- def test_arrow_batch_slicing(self):
- m, n = 100000, 10000
-
- df1 = self.spark.range(m).select((sf.col("id") % 2).alias("key"),
sf.col("id").alias("v"))
- cols = {f"col_{i}": sf.col("v") + i for i in range(10)}
- df1 = df1.withColumns(cols)
-
- df2 = self.spark.range(n).select((sf.col("id") % 4).alias("key"),
sf.col("id").alias("v"))
- cols = {f"col_{i}": sf.col("v") + i for i in range(20)}
- df2 = df2.withColumns(cols)
-
- def summarize(key, left, right):
- assert len(left) == m / 2 or len(left) == 0, len(left)
- assert len(right) == n / 4, len(right)
- return pd.DataFrame(
- {
- "key": [key[0]],
- "left_rows": [len(left)],
- "left_columns": [len(left.columns)],
- "right_rows": [len(right)],
- "right_columns": [len(right.columns)],
- }
- )
-
- schema = "key long, left_rows long, left_columns long, right_rows
long, right_columns long"
-
- expected = [
- Row(key=0, left_rows=m / 2, left_columns=12, right_rows=n / 4,
right_columns=22),
- Row(key=1, left_rows=m / 2, left_columns=12, right_rows=n / 4,
right_columns=22),
- Row(key=2, left_rows=0, left_columns=12, right_rows=n / 4,
right_columns=22),
- Row(key=3, left_rows=0, left_columns=12, right_rows=n / 4,
right_columns=22),
- ]
-
- for maxRecords, maxBytes in [(1000, 2**31 - 1), (0, 4096), (1000,
4096)]:
- with self.subTest(maxRecords=maxRecords, maxBytes=maxBytes):
- with self.sql_conf(
- {
- "spark.sql.execution.arrow.maxRecordsPerBatch":
maxRecords,
- "spark.sql.execution.arrow.maxBytesPerBatch": maxBytes,
- }
- ):
- result = (
- df1.groupby("key")
- .cogroup(df2.groupby("key"))
- .applyInPandas(summarize, schema=schema)
- .sort("key")
- .collect()
- )
-
- self.assertEqual(expected, result)
-
def test_negative_and_zero_batch_size(self):
for batch_size in [0, -1]:
with
self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": batch_size}):
CogroupedApplyInPandasTestsMixin.test_with_key_right(self)
- @unittest.skipIf(is_remote_only(), "Requires JVM access")
- def test_cogroup_apply_in_pandas_with_logging(self):
- import pandas as pd
-
- def func_with_logging(left_pdf, right_pdf):
- assert isinstance(left_pdf, pd.DataFrame)
- assert isinstance(right_pdf, pd.DataFrame)
- logger = logging.getLogger("test_pandas_cogrouped_map")
- logger.warning(
- f"pandas cogrouped map: {dict(v1=list(left_pdf['v1']),
v2=list(right_pdf['v2']))}"
- )
- return pd.merge(left_pdf, right_pdf, on=["id"])
-
- left_df = self.spark.createDataFrame([(1, 10), (2, 20), (1, 30)],
["id", "v1"])
- right_df = self.spark.createDataFrame([(1, 100), (2, 200), (1, 300)],
["id", "v2"])
-
- grouped_left = left_df.groupBy("id")
- grouped_right = right_df.groupBy("id")
- cogrouped_df = grouped_left.cogroup(grouped_right)
-
- with self.sql_conf({"spark.sql.pyspark.worker.logging.enabled":
"true"}):
- assertDataFrameEqual(
- cogrouped_df.applyInPandas(func_with_logging, "id long, v1
long, v2 long"),
- [Row(id=1, v1=v1, v2=v2) for v1 in [10, 30] for v2 in [100,
300]]
- + [Row(id=2, v1=20, v2=200)],
- )
-
- logs = self.spark.tvf.python_worker_logs()
-
- assertDataFrameEqual(
- logs.select("level", "msg", "context", "logger"),
- [
- Row(
- level="WARNING",
- msg=f"pandas cogrouped map: {dict(v1=v1, v2=v2)}",
- context={"func_name": func_with_logging.__name__},
- logger="test_pandas_cogrouped_map",
- )
- for v1, v2 in [([10, 30], [100, 300]), ([20], [200])]
- ],
- )
-
class CogroupedApplyInPandasTests(CogroupedApplyInPandasTestsMixin,
ReusedSQLTestCase):
pass
diff --git a/python/pyspark/sql/tests/pandas/test_pandas_cogrouped_map_misc.py
b/python/pyspark/sql/tests/pandas/test_pandas_cogrouped_map_misc.py
new file mode 100644
index 000000000000..ed03cedb0971
--- /dev/null
+++ b/python/pyspark/sql/tests/pandas/test_pandas_cogrouped_map_misc.py
@@ -0,0 +1,239 @@
+#
+# 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 logging
+
+from pyspark.sql import functions as sf
+from pyspark.sql.types import Row
+from pyspark.sql.window import Window
+from pyspark.testing.sqlutils import ReusedSQLTestCase
+from pyspark.testing.utils import (
+ assertDataFrameEqual,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+from pyspark.util import is_remote_only
+
+if have_pandas:
+ import pandas as pd
+
+if have_pyarrow:
+ import pyarrow as pa # noqa: F401
+
+
[email protected](
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message,
+)
+class CogroupedApplyInPandasMiscTestsMixin:
+ def test_with_window_function(self):
+ # SPARK-42168: a window function with same partition keys but
differing key order
+ ids = 2
+ days = 100
+ vals = 10000
+ parts = 10
+
+ id_df = self.spark.range(ids)
+ day_df = self.spark.range(days).withColumnRenamed("id", "day")
+ vals_df = self.spark.range(vals).withColumnRenamed("id", "value")
+ df = id_df.join(day_df).join(vals_df)
+
+ left_df = df.withColumnRenamed("value",
"left").repartition(parts).cache()
+ # SPARK-42132: this bug requires us to alias all columns from df here
+ right_df = (
+ df.select(
+ sf.col("id").alias("id"), sf.col("day").alias("day"),
sf.col("value").alias("right")
+ )
+ .repartition(parts)
+ .cache()
+ )
+
+ # note the column order is different to the groupBy("id", "day")
column order below
+ window = Window.partitionBy("day", "id")
+
+ left_grouped_df = left_df.groupBy("id", "day")
+ right_grouped_df = right_df.withColumn(
+ "day_sum", sf.sum(sf.col("day")).over(window)
+ ).groupBy("id", "day")
+
+ def cogroup(left: "pd.DataFrame", right: "pd.DataFrame") ->
"pd.DataFrame":
+ return pd.DataFrame(
+ [
+ {
+ "id": (
+ left["id"][0]
+ if not left.empty
+ else (right["id"][0] if not right.empty else None)
+ ),
+ "day": (
+ left["day"][0]
+ if not left.empty
+ else (right["day"][0] if not right.empty else None)
+ ),
+ "lefts": len(left.index),
+ "rights": len(right.index),
+ }
+ ]
+ )
+
+ df = left_grouped_df.cogroup(right_grouped_df).applyInPandas(
+ cogroup, schema="id long, day long, lefts integer, rights integer"
+ )
+
+ actual = df.orderBy("id", "day").take(days)
+ self.assertEqual(actual, [Row(0, day, vals, vals) for day in
range(days)])
+
+ def test_with_local_data(self):
+ df1 = self.spark.createDataFrame(
+ [(1, 1.0, "a"), (2, 2.0, "b"), (1, 3.0, "c"), (2, 4.0, "d")],
("id", "v1", "v2")
+ )
+ df2 = self.spark.createDataFrame([(1, "x"), (2, "y"), (1, "z")],
("id", "v3"))
+
+ def summarize(left, right):
+ return pd.DataFrame(
+ {
+ "left_rows": [len(left)],
+ "left_columns": [len(left.columns)],
+ "right_rows": [len(right)],
+ "right_columns": [len(right.columns)],
+ }
+ )
+
+ df = (
+ df1.groupby("id")
+ .cogroup(df2.groupby("id"))
+ .applyInPandas(
+ summarize,
+ schema="left_rows long, left_columns long, right_rows long,
right_columns long",
+ )
+ )
+
+ self.assertEqual(
+ df._show_string(),
+ "+---------+------------+----------+-------------+\n"
+ "|left_rows|left_columns|right_rows|right_columns|\n"
+ "+---------+------------+----------+-------------+\n"
+ "| 2| 3| 2| 2|\n"
+ "| 2| 3| 1| 2|\n"
+ "+---------+------------+----------+-------------+\n",
+ )
+
+ def test_arrow_batch_slicing(self):
+ m, n = 100000, 10000
+
+ df1 = self.spark.range(m).select((sf.col("id") % 2).alias("key"),
sf.col("id").alias("v"))
+ cols = {f"col_{i}": sf.col("v") + i for i in range(10)}
+ df1 = df1.withColumns(cols)
+
+ df2 = self.spark.range(n).select((sf.col("id") % 4).alias("key"),
sf.col("id").alias("v"))
+ cols = {f"col_{i}": sf.col("v") + i for i in range(20)}
+ df2 = df2.withColumns(cols)
+
+ def summarize(key, left, right):
+ assert len(left) == m / 2 or len(left) == 0, len(left)
+ assert len(right) == n / 4, len(right)
+ return pd.DataFrame(
+ {
+ "key": [key[0]],
+ "left_rows": [len(left)],
+ "left_columns": [len(left.columns)],
+ "right_rows": [len(right)],
+ "right_columns": [len(right.columns)],
+ }
+ )
+
+ schema = "key long, left_rows long, left_columns long, right_rows
long, right_columns long"
+
+ expected = [
+ Row(key=0, left_rows=m / 2, left_columns=12, right_rows=n / 4,
right_columns=22),
+ Row(key=1, left_rows=m / 2, left_columns=12, right_rows=n / 4,
right_columns=22),
+ Row(key=2, left_rows=0, left_columns=12, right_rows=n / 4,
right_columns=22),
+ Row(key=3, left_rows=0, left_columns=12, right_rows=n / 4,
right_columns=22),
+ ]
+
+ for maxRecords, maxBytes in [(1000, 2**31 - 1), (0, 4096), (1000,
4096)]:
+ with self.subTest(maxRecords=maxRecords, maxBytes=maxBytes):
+ with self.sql_conf(
+ {
+ "spark.sql.execution.arrow.maxRecordsPerBatch":
maxRecords,
+ "spark.sql.execution.arrow.maxBytesPerBatch": maxBytes,
+ }
+ ):
+ result = (
+ df1.groupby("key")
+ .cogroup(df2.groupby("key"))
+ .applyInPandas(summarize, schema=schema)
+ .sort("key")
+ .collect()
+ )
+
+ self.assertEqual(expected, result)
+
+ @unittest.skipIf(is_remote_only(), "Requires JVM access")
+ def test_cogroup_apply_in_pandas_with_logging(self):
+ import pandas as pd
+
+ def func_with_logging(left_pdf, right_pdf):
+ assert isinstance(left_pdf, pd.DataFrame)
+ assert isinstance(right_pdf, pd.DataFrame)
+ logger = logging.getLogger("test_pandas_cogrouped_map")
+ logger.warning(
+ f"pandas cogrouped map: {dict(v1=list(left_pdf['v1']),
v2=list(right_pdf['v2']))}"
+ )
+ return pd.merge(left_pdf, right_pdf, on=["id"])
+
+ left_df = self.spark.createDataFrame([(1, 10), (2, 20), (1, 30)],
["id", "v1"])
+ right_df = self.spark.createDataFrame([(1, 100), (2, 200), (1, 300)],
["id", "v2"])
+
+ grouped_left = left_df.groupBy("id")
+ grouped_right = right_df.groupBy("id")
+ cogrouped_df = grouped_left.cogroup(grouped_right)
+
+ with self.sql_conf({"spark.sql.pyspark.worker.logging.enabled":
"true"}):
+ assertDataFrameEqual(
+ cogrouped_df.applyInPandas(func_with_logging, "id long, v1
long, v2 long"),
+ [Row(id=1, v1=v1, v2=v2) for v1 in [10, 30] for v2 in [100,
300]]
+ + [Row(id=2, v1=20, v2=200)],
+ )
+
+ logs = self.spark.tvf.python_worker_logs()
+
+ assertDataFrameEqual(
+ logs.select("level", "msg", "context", "logger"),
+ [
+ Row(
+ level="WARNING",
+ msg=f"pandas cogrouped map: {dict(v1=v1, v2=v2)}",
+ context={"func_name": func_with_logging.__name__},
+ logger="test_pandas_cogrouped_map",
+ )
+ for v1, v2 in [([10, 30], [100, 300]), ([20], [200])]
+ ],
+ )
+
+
+class CogroupedApplyInPandasMiscTests(CogroupedApplyInPandasMiscTestsMixin,
ReusedSQLTestCase):
+ pass
+
+
+if __name__ == "__main__":
+ from pyspark.testing import main
+
+ main()
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
