#
# 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 numpy as np
import pandas as pd

from pyspark import pandas as ps
from pyspark.loose_version import LooseVersion
from pyspark.pandas.config import option_context
from pyspark.testing.pandasutils import PandasOnSparkTestCase
from pyspark.testing.sqlutils import SQLTestUtils


class GroupbyApplyFuncMixin:
    def test_apply(self):
        pdf = pd.DataFrame(
            {"a": [1, 2, 3, 4, 5, 6], "b": [1, 1, 2, 3, 5, 8], "c": [1, 4, 9, 16, 25, 36]},
            columns=["a", "b", "c"],
        )
        psdf = ps.from_pandas(pdf)

        if LooseVersion(pd.__version__) < "3.0.0":
            for include_groups in [True, False]:
                with self.subTest(include_groups=include_groups):
                    self._check_apply(psdf, pdf, include_groups)
        else:
            self._check_apply(psdf, pdf, include_groups=False)
            with self.assertRaises(ValueError):
                psdf.groupby("b").apply(lambda x: x + x.min(), include_groups=True)

    def _check_apply(self, psdf, pdf, include_groups):
        self.assert_eq(
            psdf.groupby("b")
            .apply(lambda x: x + x.min(), include_groups=include_groups)
            .sort_index(),
            pdf.groupby("b")
            .apply(lambda x: x + x.min(), include_groups=include_groups)
            .sort_index(),
        )
        self.assert_eq(
            psdf.groupby("b").apply(len, include_groups=include_groups).sort_index(),
            pdf.groupby("b").apply(len, include_groups=include_groups).sort_index(),
        )
        self.assert_eq(
            psdf.groupby("b")["a"]
            .apply(lambda x, y, z: x + x.min() + y * z, 10, z=20, include_groups=include_groups)
            .sort_index(),
            pdf.groupby("b")["a"]
            .apply(lambda x, y, z: x + x.min() + y * z, 10, z=20, include_groups=include_groups)
            .sort_index(),
        )
        self.assert_eq(
            psdf.groupby("b")[["a"]]
            .apply(lambda x: x + x.min(), include_groups=include_groups)
            .sort_index(),
            pdf.groupby("b")[["a"]]
            .apply(lambda x: x + x.min(), include_groups=include_groups)
            .sort_index(),
        )
        self.assert_eq(
            psdf.groupby(["a", "b"])
            .apply(lambda x, y, z: x + x.min() + y + z, 1, z=2, include_groups=include_groups)
            .sort_index(),
            pdf.groupby(["a", "b"])
            .apply(lambda x, y, z: x + x.min() + y + z, 1, z=2, include_groups=include_groups)
            .sort_index(),
        )
        self.assert_eq(
            psdf.groupby(["b"])["c"].apply(lambda x: 1, include_groups=include_groups).sort_index(),
            pdf.groupby(["b"])["c"].apply(lambda x: 1, include_groups=include_groups).sort_index(),
        )
        self.assert_eq(
            psdf.groupby(["b"])["c"].apply(len, include_groups=include_groups).sort_index(),
            pdf.groupby(["b"])["c"].apply(len, include_groups=include_groups).sort_index(),
        )
        self.assert_eq(
            psdf.groupby(psdf.b // 5)
            .apply(lambda x: x + x.min(), include_groups=include_groups)
            .sort_index(),
            pdf.groupby(pdf.b // 5)
            .apply(lambda x: x + x.min(), include_groups=include_groups)
            .sort_index(),
            almost=True,
        )
        self.assert_eq(
            psdf.groupby(psdf.b // 5)["a"]
            .apply(lambda x: x + x.min(), include_groups=include_groups)
            .sort_index(),
            pdf.groupby(pdf.b // 5)["a"]
            .apply(lambda x: x + x.min(), include_groups=include_groups)
            .sort_index(),
            almost=True,
        )
        self.assert_eq(
            psdf.groupby(psdf.b // 5)[["a"]]
            .apply(lambda x: x + x.min(), include_groups=include_groups)
            .sort_index(),
            pdf.groupby(pdf.b // 5)[["a"]]
            .apply(lambda x: x + x.min(), include_groups=include_groups)
            .sort_index(),
            almost=True,
        )
        self.assert_eq(
            psdf.groupby(psdf.b // 5)[["a"]].apply(len, include_groups=include_groups).sort_index(),
            pdf.groupby(pdf.b // 5)[["a"]].apply(len, include_groups=include_groups).sort_index(),
            almost=True,
        )
        self.assert_eq(
            psdf.a.rename()
            .groupby(psdf.b)
            .apply(lambda x: x + x.min(), include_groups=include_groups)
            .sort_index(),
            pdf.a.rename()
            .groupby(pdf.b)
            .apply(lambda x: x + x.min(), include_groups=include_groups)
            .sort_index(),
        )
        self.assert_eq(
            psdf.a.groupby(psdf.b.rename())
            .apply(lambda x: x + x.min(), include_groups=include_groups)
            .sort_index(),
            pdf.a.groupby(pdf.b.rename())
            .apply(lambda x: x + x.min(), include_groups=include_groups)
            .sort_index(),
        )
        self.assert_eq(
            psdf.a.rename()
            .groupby(psdf.b.rename())
            .apply(lambda x: x + x.min(), include_groups=include_groups)
            .sort_index(),
            pdf.a.rename()
            .groupby(pdf.b.rename())
            .apply(lambda x: x + x.min(), include_groups=include_groups)
            .sort_index(),
        )

        with self.assertRaisesRegex(TypeError, "int object is not callable"):
            psdf.groupby("b").apply(1, include_groups=include_groups)

    def test_apply_with_multi_index_columns(self):
        pdf = pd.DataFrame(
            {"a": [1, 2, 3, 4, 5, 6], "b": [1, 1, 2, 3, 5, 8], "c": [1, 4, 9, 16, 25, 36]},
            columns=["a", "b", "c"],
        )
        psdf = ps.from_pandas(pdf)

        # multi-index columns
        columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")])
        pdf.columns = columns
        psdf.columns = columns

        if LooseVersion(pd.__version__) < "3.0.0":
            for include_groups in [True, False]:
                with self.subTest(include_groups=include_groups):
                    self._check_apply_with_multi_index_columns(psdf, pdf, include_groups)
        else:
            self._check_apply_with_multi_index_columns(psdf, pdf, include_groups=False)
            with self.assertRaises(ValueError):
                psdf.groupby(("x", "b")).apply(lambda x: x + x.min(), include_groups=True)

    def _check_apply_with_multi_index_columns(self, psdf, pdf, include_groups):
        self.assert_eq(
            psdf.groupby(("x", "b")).apply(lambda x: 1, include_groups=include_groups).sort_index(),
            pdf.groupby(("x", "b")).apply(lambda x: 1, include_groups=include_groups).sort_index(),
        )
        self.assert_eq(
            psdf.groupby([("x", "a"), ("x", "b")])
            .apply(lambda x: x + x.min(), include_groups=include_groups)
            .sort_index(),
            pdf.groupby([("x", "a"), ("x", "b")])
            .apply(lambda x: x + x.min(), include_groups=include_groups)
            .sort_index(),
        )
        self.assert_eq(
            psdf.groupby(("x", "b")).apply(len, include_groups=include_groups).sort_index(),
            pdf.groupby(("x", "b")).apply(len, include_groups=include_groups).sort_index(),
        )
        self.assert_eq(
            psdf.groupby([("x", "a"), ("x", "b")])
            .apply(len, include_groups=include_groups)
            .sort_index(),
            pdf.groupby([("x", "a"), ("x", "b")])
            .apply(len, include_groups=include_groups)
            .sort_index(),
        )

    def test_apply_without_shortcut(self):
        with option_context("compute.shortcut_limit", 0):
            self.test_apply()

    def test_apply_with_multi_index_columns_without_shortcut(self):
        with option_context("compute.shortcut_limit", 0):
            self.test_apply_with_multi_index_columns()

    def test_apply_with_type_hint(self):
        pdf = pd.DataFrame(
            {"a": [1, 2, 3, 4, 5, 6], "b": [1, 1, 2, 3, 5, 8], "c": [1, 4, 9, 16, 25, 36]},
            columns=["a", "b", "c"],
        )
        psdf = ps.from_pandas(pdf)

        if LooseVersion(pd.__version__) < "3.0.0":
            for include_groups in [True, False]:
                with self.subTest(include_groups=include_groups):
                    self._check_apply_with_type_hint(psdf, pdf, include_groups)
        else:
            self._check_apply_with_type_hint(psdf, pdf, include_groups=False)

    def _check_apply_with_type_hint(self, psdf, pdf, include_groups):
        if include_groups:

            def add_max1(x) -> ps.DataFrame[int, int, int]:
                return x + x.min()

        else:

            def add_max1(x) -> ps.DataFrame[int, int]:
                return x + x.min()

        # Type hints set the default column names, and we use default index for
        # pandas API on Spark. Here we ignore both diff.
        actual = psdf.groupby("b").apply(add_max1, include_groups=include_groups).sort_index()
        expected = pdf.groupby("b").apply(add_max1, include_groups=include_groups).sort_index()

        if include_groups:
            self.assert_eq(sorted(actual["c0"].to_numpy()), sorted(expected["a"].to_numpy()))
            self.assert_eq(sorted(actual["c1"].to_numpy()), sorted(expected["b"].to_numpy()))
            self.assert_eq(sorted(actual["c2"].to_numpy()), sorted(expected["c"].to_numpy()))
        else:
            self.assert_eq(sorted(actual["c0"].to_numpy()), sorted(expected["a"].to_numpy()))
            self.assert_eq(sorted(actual["c1"].to_numpy()), sorted(expected["c"].to_numpy()))

        if include_groups:

            def add_max2(
                x,
            ) -> ps.DataFrame[slice("a", int), slice("b", int), slice("c", int)]:
                return x + x.min()

        else:

            def add_max2(
                x,
            ) -> ps.DataFrame[slice("a", int), slice("c", int)]:
                return x + x.min()

        actual = psdf.groupby("b").apply(add_max2, include_groups=include_groups).sort_index()
        expected = pdf.groupby("b").apply(add_max2, include_groups=include_groups).sort_index()

        if include_groups:
            self.assert_eq(sorted(actual["a"].to_numpy()), sorted(expected["a"].to_numpy()))
            self.assert_eq(sorted(actual["b"].to_numpy()), sorted(expected["b"].to_numpy()))
            self.assert_eq(sorted(actual["c"].to_numpy()), sorted(expected["c"].to_numpy()))
        else:
            self.assert_eq(sorted(actual["a"].to_numpy()), sorted(expected["a"].to_numpy()))
            self.assert_eq(sorted(actual["c"].to_numpy()), sorted(expected["c"].to_numpy()))

    def test_apply_negative(self):
        def func(_) -> ps.Series[int]:
            return pd.Series([1])

        with self.assertRaisesRegex(TypeError, "Series as a return type hint at frame groupby"):
            ps.range(10).groupby("id").apply(func)

    def test_apply_with_new_dataframe(self):
        pdf = pd.DataFrame(
            {"timestamp": [0.0, 0.5, 1.0, 0.0, 0.5], "car_id": ["A", "A", "A", "B", "B"]}
        )
        psdf = ps.from_pandas(pdf)

        self.assert_eq(
            psdf.groupby("car_id").apply(lambda _: pd.DataFrame({"column": [0.0]})).sort_index(),
            pdf.groupby("car_id").apply(lambda _: pd.DataFrame({"column": [0.0]})).sort_index(),
        )

        self.assert_eq(
            psdf.groupby("car_id")
            .apply(lambda df: pd.DataFrame({"mean": [df["timestamp"].mean()]}))
            .sort_index(),
            pdf.groupby("car_id")
            .apply(lambda df: pd.DataFrame({"mean": [df["timestamp"].mean()]}))
            .sort_index(),
        )

        # dataframe with 1000+ records
        pdf = pd.DataFrame(
            {
                "timestamp": [0.0, 0.5, 1.0, 0.0, 0.5] * 300,
                "car_id": ["A", "A", "A", "B", "B"] * 300,
            }
        )
        psdf = ps.from_pandas(pdf)

        self.assert_eq(
            psdf.groupby("car_id").apply(lambda _: pd.DataFrame({"column": [0.0]})).sort_index(),
            pdf.groupby("car_id").apply(lambda _: pd.DataFrame({"column": [0.0]})).sort_index(),
        )

        self.assert_eq(
            psdf.groupby("car_id")
            .apply(lambda df: pd.DataFrame({"mean": [df["timestamp"].mean()]}))
            .sort_index(),
            pdf.groupby("car_id")
            .apply(lambda df: pd.DataFrame({"mean": [df["timestamp"].mean()]}))
            .sort_index(),
        )

    def test_apply_infer_schema_without_shortcut(self):
        # SPARK-39054: Ensure infer schema accuracy in GroupBy.apply
        with option_context("compute.shortcut_limit", 0):
            dfs = (
                {"timestamp": [0.0], "car_id": ["A"]},
                {"timestamp": [0.0, 0.0], "car_id": ["A", "A"]},
            )
            func = lambda _: pd.DataFrame({"column": [0.0]})  # noqa: E731
            for df in dfs:
                pdf = pd.DataFrame(df)
                psdf = ps.from_pandas(pdf)
                self.assert_eq(
                    psdf.groupby("car_id").apply(func).sort_index(),
                    pdf.groupby("car_id").apply(func).sort_index(),
                )

    def test_apply_with_new_dataframe_without_shortcut(self):
        with option_context("compute.shortcut_limit", 0):
            self.test_apply_with_new_dataframe()

    def test_apply_key_handling(self):
        pdf = pd.DataFrame(
            {"d": [1.0, 1.0, 1.0, 2.0, 2.0, 2.0], "v": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]}
        )
        psdf = ps.from_pandas(pdf)

        if LooseVersion(pd.__version__) < "3.0.0":
            sum_f = sum
        else:
            sum_f = np.sum

        self.assert_eq(
            psdf.groupby("d").apply(sum_f).sort_index(), pdf.groupby("d").apply(sum_f).sort_index()
        )

    def test_apply_key_handling_without_shortcut(self):
        with ps.option_context("compute.shortcut_limit", 0):
            self.test_apply_key_handling()

    def test_apply_with_side_effect(self):
        pdf = pd.DataFrame(
            {"d": [1.0, 1.0, 1.0, 2.0, 2.0, 2.0], "v": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]}
        )
        psdf = ps.from_pandas(pdf)

        if LooseVersion(pd.__version__) < "3.0.0":
            for include_groups in [True, False]:
                with self.subTest(include_groups=include_groups):
                    self._check_apply_with_side_effect(psdf, pdf, include_groups)
        else:
            self._check_apply_with_side_effect(psdf, pdf, include_groups=False)

    def _check_apply_with_side_effect(self, psdf, pdf, include_groups):
        acc = ps.utils.default_session().sparkContext.accumulator(0)

        if LooseVersion(pd.__version__) < "3.0.0":
            sum_f = sum
        else:
            sum_f = np.sum

        if include_groups:

            def sum_with_acc_frame(x) -> ps.DataFrame[np.float64, np.float64]:
                nonlocal acc
                acc += 1
                return np.sum(x)

        else:
            if LooseVersion(pd.__version__) < "3.0.0":
                ret_type = ps.DataFrame[np.float64]
            else:
                ret_type = np.float64

            def sum_with_acc_frame(x) -> ret_type:
                nonlocal acc
                acc += 1
                return np.sum(x)

        actual = psdf.groupby("d").apply(sum_with_acc_frame, include_groups=include_groups)
        if LooseVersion(pd.__version__) < "3.0.0":
            actual.columns = ["d", "v"] if include_groups else ["v"]
        else:
            actual = actual.rename()
        self.assert_eq(
            actual._to_pandas().sort_index(),
            pdf.groupby("d")
            .apply(sum_f, include_groups=include_groups)
            .sort_index()
            .reset_index(drop=True),
        )
        self.assert_eq(acc.value, 2)

        def sum_with_acc_series(x) -> np.float64:
            nonlocal acc
            acc += 1
            return np.sum(x)

        self.assert_eq(
            psdf.groupby("d")["v"]
            .apply(sum_with_acc_series, include_groups=include_groups)
            ._to_pandas()
            .sort_index(),
            pdf.groupby("d")["v"]
            .apply(sum_f, include_groups=include_groups)
            .sort_index()
            .reset_index(drop=True),
        )
        self.assert_eq(acc.value, 4)

    def test_apply_with_side_effect_without_shortcut(self):
        with ps.option_context("compute.shortcut_limit", 0):
            self.test_apply_with_side_effect()

    def test_apply_return_series(self):
        # SPARK-36907: Fix DataFrameGroupBy.apply without shortcut.
        pdf = pd.DataFrame(
            {"a": [1, 2, 3, 4, 5, 6], "b": [1, 1, 2, 3, 5, 8], "c": [1, 4, 9, 16, 25, 36]},
            columns=["a", "b", "c"],
        )
        psdf = ps.from_pandas(pdf)

        if LooseVersion(pd.__version__) < "3.0.0":
            for include_groups in [True, False]:
                with self.subTest(include_groups=include_groups):
                    self._check_apply_return_series(psdf, pdf, include_groups)
        else:
            self._check_apply_return_series(psdf, pdf, include_groups=False)
            with self.assertRaises(ValueError):
                psdf.groupby("b").apply(lambda x: x + x.iloc[0], include_groups=True)

    def _check_apply_return_series(self, psdf, pdf, include_groups):
        self.assert_eq(
            psdf.groupby("b")
            .apply(lambda x: x.iloc[0], include_groups=include_groups)
            .sort_index(),
            pdf.groupby("b").apply(lambda x: x.iloc[0], include_groups=include_groups).sort_index(),
        )
        self.assert_eq(
            psdf.groupby("b").apply(lambda x: x["a"], include_groups=include_groups).sort_index(),
            pdf.groupby("b").apply(lambda x: x["a"], include_groups=include_groups).sort_index(),
        )
        self.assert_eq(
            psdf.groupby(["b", "c"])
            .apply(lambda x: x.iloc[0], include_groups=include_groups)
            .sort_index(),
            pdf.groupby(["b", "c"])
            .apply(lambda x: x.iloc[0], include_groups=include_groups)
            .sort_index(),
        )
        self.assert_eq(
            psdf.groupby(["b", "c"])
            .apply(lambda x: x["a"], include_groups=include_groups)
            .sort_index(),
            pdf.groupby(["b", "c"])
            .apply(lambda x: x["a"], include_groups=include_groups)
            .sort_index(),
        )

    def test_apply_return_series_with_multi_index_columns(self):
        # SPARK-36907: Fix DataFrameGroupBy.apply without shortcut.
        pdf = pd.DataFrame(
            {"a": [1, 2, 3, 4, 5, 6], "b": [1, 1, 2, 3, 5, 8], "c": [1, 4, 9, 16, 25, 36]},
            columns=["a", "b", "c"],
        )
        psdf = ps.from_pandas(pdf)

        # multi-index columns
        columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")])
        pdf.columns = columns
        psdf.columns = columns

        if LooseVersion(pd.__version__) < "3.0.0":
            for include_groups in [True, False]:
                with self.subTest(include_groups=include_groups):
                    self._check_apply_return_series_with_multi_index_columns(
                        psdf, pdf, include_groups
                    )
        else:
            self._check_apply_return_series_with_multi_index_columns(
                psdf, pdf, include_groups=False
            )
            with self.assertRaises(ValueError):
                psdf.groupby(("x", "b")).apply(lambda x: x + x.iloc[0], include_groups=True)

    def _check_apply_return_series_with_multi_index_columns(self, psdf, pdf, include_groups):
        self.assert_eq(
            psdf.groupby(("x", "b"))
            .apply(lambda x: x.iloc[0], include_groups=include_groups)
            .sort_index(),
            pdf.groupby(("x", "b"))
            .apply(lambda x: x.iloc[0], include_groups=include_groups)
            .sort_index(),
        )
        self.assert_eq(
            psdf.groupby(("x", "b"))
            .apply(lambda x: x[("x", "a")], include_groups=include_groups)
            .sort_index(),
            pdf.groupby(("x", "b"))
            .apply(lambda x: x[("x", "a")], include_groups=include_groups)
            .sort_index(),
        )
        self.assert_eq(
            psdf.groupby([("x", "b"), ("y", "c")])
            .apply(lambda x: x.iloc[0], include_groups=include_groups)
            .sort_index(),
            pdf.groupby([("x", "b"), ("y", "c")])
            .apply(lambda x: x.iloc[0], include_groups=include_groups)
            .sort_index(),
        )
        self.assert_eq(
            psdf.groupby([("x", "b"), ("y", "c")])
            .apply(lambda x: x[("x", "a")], include_groups=include_groups)
            .sort_index(),
            pdf.groupby([("x", "b"), ("y", "c")])
            .apply(lambda x: x[("x", "a")], include_groups=include_groups)
            .sort_index(),
        )

    def test_apply_return_series_without_shortcut(self):
        # SPARK-36907: Fix DataFrameGroupBy.apply without shortcut.
        with ps.option_context("compute.shortcut_limit", 2):
            self.test_apply_return_series()

    def test_apply_return_series_with_multi_index_columns_without_shortcut(self):
        with ps.option_context("compute.shortcut_limit", 2):
            self.test_apply_return_series_with_multi_index_columns()

    def test_apply_explicitly_infer(self):
        # SPARK-39317
        from pyspark.pandas.utils import SPARK_CONF_ARROW_ENABLED

        def plus_min(x):
            return x + x.min()

        with self.sql_conf({SPARK_CONF_ARROW_ENABLED: False}):
            df = ps.DataFrame({"A": ["a", "a", "b"], "B": [1, 2, 3]}, columns=["A", "B"])
            g = df.groupby("A")
            g.apply(plus_min).sort_index()

    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"],
        )
        psdf = ps.from_pandas(pdf)
        self.assert_eq(
            psdf.groupby("b").transform(lambda x: x + x.min()).sort_index(),
            pdf.groupby("b").transform(lambda x: x + x.min()).sort_index(),
        )
        self.assert_eq(
            psdf.groupby("b")["a"].transform(lambda x: x + x.min()).sort_index(),
            pdf.groupby("b")["a"].transform(lambda x: x + x.min()).sort_index(),
        )
        self.assert_eq(
            psdf.groupby("b")[["a"]].transform(lambda x: x + x.min()).sort_index(),
            pdf.groupby("b")[["a"]].transform(lambda x: x + x.min()).sort_index(),
        )
        self.assert_eq(
            psdf.groupby(["a", "b"]).transform(lambda x: x + x.min()).sort_index(),
            pdf.groupby(["a", "b"]).transform(lambda x: x + x.min()).sort_index(),
        )
        self.assert_eq(
            psdf.groupby(["b"])["c"].transform(lambda x: x + x.min()).sort_index(),
            pdf.groupby(["b"])["c"].transform(lambda x: x + x.min()).sort_index(),
        )
        self.assert_eq(
            psdf.groupby(psdf.b // 5).transform(lambda x: x + x.min()).sort_index(),
            pdf.groupby(pdf.b // 5).transform(lambda x: x + x.min()).sort_index(),
        )
        self.assert_eq(
            psdf.groupby(psdf.b // 5)["a"].transform(lambda x: x + x.min()).sort_index(),
            pdf.groupby(pdf.b // 5)["a"].transform(lambda x: x + x.min()).sort_index(),
        )
        self.assert_eq(
            psdf.groupby(psdf.b // 5)[["a"]].transform(lambda x: x + x.min()).sort_index(),
            pdf.groupby(pdf.b // 5)[["a"]].transform(lambda x: x + x.min()).sort_index(),
        )
        self.assert_eq(
            psdf.a.rename().groupby(psdf.b).transform(lambda x: x + x.min()).sort_index(),
            pdf.a.rename().groupby(pdf.b).transform(lambda x: x + x.min()).sort_index(),
        )
        self.assert_eq(
            psdf.a.groupby(psdf.b.rename()).transform(lambda x: x + x.min()).sort_index(),
            pdf.a.groupby(pdf.b.rename()).transform(lambda x: x + x.min()).sort_index(),
        )
        self.assert_eq(
            psdf.a.rename().groupby(psdf.b.rename()).transform(lambda x: x + x.min()).sort_index(),
            pdf.a.rename().groupby(pdf.b.rename()).transform(lambda x: x + x.min()).sort_index(),
        )
        with self.assertRaisesRegex(TypeError, "str object is not callable"):
            psdf.groupby("a").transform("sum")

        def udf(col) -> int:
            return col + 10

        with self.assertRaisesRegex(
            TypeError,
            "Expected the return type of this function to be of Series type, "
            "but found type ScalarType\\[LongType\\(\\)\\]",
        ):
            psdf.groupby("a").transform(udf)

        # multi-index columns
        columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")])
        pdf.columns = columns
        psdf.columns = columns

        self.assert_eq(
            psdf.groupby(("x", "b")).transform(lambda x: x + x.min()).sort_index(),
            pdf.groupby(("x", "b")).transform(lambda x: x + x.min()).sort_index(),
        )
        self.assert_eq(
            psdf.groupby([("x", "a"), ("x", "b")]).transform(lambda x: x + x.min()).sort_index(),
            pdf.groupby([("x", "a"), ("x", "b")]).transform(lambda x: x + x.min()).sort_index(),
        )

    def test_transform_without_shortcut(self):
        with option_context("compute.shortcut_limit", 0):
            self.test_transform()

    def test_filter(self):
        pdf = pd.DataFrame(
            {"a": [1, 2, 3, 4, 5, 6], "b": [1, 1, 2, 3, 5, 8], "c": [1, 4, 9, 16, 25, 36]},
            columns=["a", "b", "c"],
        )
        psdf = ps.from_pandas(pdf)

        self.assert_eq(
            psdf.groupby("b").filter(lambda x: any(x.a == 2)).sort_index(),
            pdf.groupby("b").filter(lambda x: any(x.a == 2)).sort_index(),
        )
        self.assert_eq(
            psdf.groupby("b")["a"].filter(lambda x: any(x == 2)).sort_index(),
            pdf.groupby("b")["a"].filter(lambda x: any(x == 2)).sort_index(),
        )
        self.assert_eq(
            psdf.groupby("b")[["a"]].filter(lambda x: any(x.a == 2)).sort_index(),
            pdf.groupby("b")[["a"]].filter(lambda x: any(x.a == 2)).sort_index(),
        )
        self.assert_eq(
            psdf.groupby(["a", "b"]).filter(lambda x: any(x.a == 2)).sort_index(),
            pdf.groupby(["a", "b"]).filter(lambda x: any(x.a == 2)).sort_index(),
        )
        self.assert_eq(
            psdf.groupby(psdf["b"] // 5).filter(lambda x: any(x.a == 2)).sort_index(),
            pdf.groupby(pdf["b"] // 5).filter(lambda x: any(x.a == 2)).sort_index(),
        )
        self.assert_eq(
            psdf.groupby(psdf["b"] // 5)["a"].filter(lambda x: any(x == 2)).sort_index(),
            pdf.groupby(pdf["b"] // 5)["a"].filter(lambda x: any(x == 2)).sort_index(),
        )
        self.assert_eq(
            psdf.groupby(psdf["b"] // 5)[["a"]].filter(lambda x: any(x.a == 2)).sort_index(),
            pdf.groupby(pdf["b"] // 5)[["a"]].filter(lambda x: any(x.a == 2)).sort_index(),
        )
        self.assert_eq(
            psdf.a.rename().groupby(psdf.b).filter(lambda x: any(x == 2)).sort_index(),
            pdf.a.rename().groupby(pdf.b).filter(lambda x: any(x == 2)).sort_index(),
        )
        self.assert_eq(
            psdf.a.groupby(psdf.b.rename()).filter(lambda x: any(x == 2)).sort_index(),
            pdf.a.groupby(pdf.b.rename()).filter(lambda x: any(x == 2)).sort_index(),
        )
        self.assert_eq(
            psdf.a.rename().groupby(psdf.b.rename()).filter(lambda x: any(x == 2)).sort_index(),
            pdf.a.rename().groupby(pdf.b.rename()).filter(lambda x: any(x == 2)).sort_index(),
        )

        with self.assertRaisesRegex(TypeError, "int object is not callable"):
            psdf.groupby("b").filter(1)

        # multi-index columns
        columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")])
        pdf.columns = columns
        psdf.columns = columns

        self.assert_eq(
            psdf.groupby(("x", "b")).filter(lambda x: any(x[("x", "a")] == 2)).sort_index(),
            pdf.groupby(("x", "b")).filter(lambda x: any(x[("x", "a")] == 2)).sort_index(),
        )
        self.assert_eq(
            psdf.groupby([("x", "a"), ("x", "b")])
            .filter(lambda x: any(x[("x", "a")] == 2))
            .sort_index(),
            pdf.groupby([("x", "a"), ("x", "b")])
            .filter(lambda x: any(x[("x", "a")] == 2))
            .sort_index(),
        )


class GroupbyApplyFuncTests(
    GroupbyApplyFuncMixin,
    PandasOnSparkTestCase,
    SQLTestUtils,
):
    pass


if __name__ == "__main__":
    from pyspark.testing import main

    main()
