# # 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 inspect from distutils.version import LooseVersion import numpy as np import pandas as pd from pyspark import pandas as ps from pyspark.pandas.exceptions import PandasNotImplementedError, DataError from pyspark.pandas.missing.groupby import ( MissingPandasLikeDataFrameGroupBy, MissingPandasLikeSeriesGroupBy, ) from pyspark.pandas.groupby import is_multi_agg_with_relabel, SeriesGroupBy from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils class GroupByTestsMixin: @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_groupby_simple(self): pdf = pd.DataFrame( { "a": [1, 2, 6, 4, 4, 6, 4, 3, 7], "b": [4, 2, 7, 3, 3, 1, 1, 1, 2], "c": [4, 2, 7, 3, None, 1, 1, 1, 2], "d": list("abcdefght"), "e": [True, False, True, False, True, False, True, False, True], }, index=[0, 1, 3, 5, 6, 8, 9, 9, 9], ) if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): # TODO(SPARK-43295): Make DataFrameGroupBy.sum support for string type columns pdf = pdf[["a", "b", "c", "e"]] psdf = ps.from_pandas(pdf) for as_index in [True, False]: if as_index: def sort(df): return df.sort_index() else: def sort(df): return df.sort_values("a").reset_index(drop=True) self.assert_eq( sort(psdf.groupby("a", as_index=as_index).sum()), sort(pdf.groupby("a", as_index=as_index).sum()), ) self.assert_eq( sort(psdf.groupby("a", as_index=as_index).b.sum()), sort(pdf.groupby("a", as_index=as_index).b.sum()), ) self.assert_eq( sort(psdf.groupby("a", as_index=as_index)["b"].sum()), sort(pdf.groupby("a", as_index=as_index)["b"].sum()), ) self.assert_eq( sort(psdf.groupby("a", as_index=as_index)[["b", "c"]].sum()), sort(pdf.groupby("a", as_index=as_index)[["b", "c"]].sum()), ) self.assert_eq( sort(psdf.groupby("a", as_index=as_index)[[]].sum()), sort(pdf.groupby("a", as_index=as_index)[[]].sum()), ) self.assert_eq( sort(psdf.groupby("a", as_index=as_index)["c"].sum()), sort(pdf.groupby("a", as_index=as_index)["c"].sum()), ) self.assert_eq( psdf.groupby("a").a.sum().sort_index(), pdf.groupby("a").a.sum().sort_index() ) self.assert_eq( psdf.groupby("a")["a"].sum().sort_index(), pdf.groupby("a")["a"].sum().sort_index() ) self.assert_eq( psdf.groupby("a")[["a"]].sum().sort_index(), pdf.groupby("a")[["a"]].sum().sort_index() ) self.assert_eq( psdf.groupby("a")[["a", "c"]].sum().sort_index(), pdf.groupby("a")[["a", "c"]].sum().sort_index(), ) self.assert_eq( psdf.a.groupby(psdf.b).sum().sort_index(), pdf.a.groupby(pdf.b).sum().sort_index() ) for axis in [0, "index"]: self.assert_eq( psdf.groupby("a", axis=axis).a.sum().sort_index(), pdf.groupby("a", axis=axis).a.sum().sort_index(), ) self.assert_eq( psdf.groupby("a", axis=axis)["a"].sum().sort_index(), pdf.groupby("a", axis=axis)["a"].sum().sort_index(), ) self.assert_eq( psdf.groupby("a", axis=axis)[["a"]].sum().sort_index(), pdf.groupby("a", axis=axis)[["a"]].sum().sort_index(), ) self.assert_eq( psdf.groupby("a", axis=axis)[["a", "c"]].sum().sort_index(), pdf.groupby("a", axis=axis)[["a", "c"]].sum().sort_index(), ) self.assert_eq( psdf.a.groupby(psdf.b, axis=axis).sum().sort_index(), pdf.a.groupby(pdf.b, axis=axis).sum().sort_index(), ) self.assertRaises(ValueError, lambda: psdf.groupby("a", as_index=False).a) self.assertRaises(ValueError, lambda: psdf.groupby("a", as_index=False)["a"]) self.assertRaises(ValueError, lambda: psdf.groupby("a", as_index=False)[["a"]]) self.assertRaises(ValueError, lambda: psdf.groupby("a", as_index=False)[["a", "c"]]) self.assertRaises(KeyError, lambda: psdf.groupby("z", as_index=False)[["a", "c"]]) self.assertRaises(KeyError, lambda: psdf.groupby(["z"], as_index=False)[["a", "c"]]) self.assertRaises(TypeError, lambda: psdf.a.groupby(psdf.b, as_index=False)) self.assertRaises(NotImplementedError, lambda: psdf.groupby("a", axis=1)) self.assertRaises(NotImplementedError, lambda: psdf.groupby("a", axis="columns")) self.assertRaises(ValueError, lambda: psdf.groupby("a", "b")) self.assertRaises(TypeError, lambda: psdf.a.groupby(psdf.a, psdf.b)) # we can't use column name/names as a parameter `by` for `SeriesGroupBy`. self.assertRaises(KeyError, lambda: psdf.a.groupby(by="a")) self.assertRaises(KeyError, lambda: psdf.a.groupby(by=["a", "b"])) self.assertRaises(KeyError, lambda: psdf.a.groupby(by=("a", "b"))) self.assertRaises(KeyError, lambda: psdf.a.groupby(by=[("a", "b")])) # we can't use DataFrame as a parameter `by` for `DataFrameGroupBy`/`SeriesGroupBy`. self.assertRaises(ValueError, lambda: psdf.groupby(psdf)) self.assertRaises(ValueError, lambda: psdf.a.groupby(psdf)) self.assertRaises(ValueError, lambda: psdf.a.groupby((psdf,))) with self.assertRaisesRegex(ValueError, "Grouper for 'list' not 1-dimensional"): psdf.groupby(by=[["a", "b"]]) # non-string names pdf = pd.DataFrame( { 10: [1, 2, 6, 4, 4, 6, 4, 3, 7], 20: [4, 2, 7, 3, 3, 1, 1, 1, 2], 30: [4, 2, 7, 3, None, 1, 1, 1, 2], 40: list("abcdefght"), }, index=[0, 1, 3, 5, 6, 8, 9, 9, 9], ) psdf = ps.from_pandas(pdf) if LooseVersion(pd.__version__) >= LooseVersion("2.0.0"): # TODO(SPARK-43295): Make DataFrameGroupBy.sum support for string type columns pdf = pdf[[10, 20, 30]] for as_index in [True, False]: if as_index: def sort(df): return df.sort_index() else: def sort(df): return df.sort_values(10).reset_index(drop=True) self.assert_eq( sort(psdf.groupby(10, as_index=as_index).sum()), sort(pdf.groupby(10, as_index=as_index).sum()), ) self.assert_eq( sort(psdf.groupby(10, as_index=as_index)[20].sum()), sort(pdf.groupby(10, as_index=as_index)[20].sum()), ) self.assert_eq( sort(psdf.groupby(10, as_index=as_index)[[20, 30]].sum()), sort(pdf.groupby(10, as_index=as_index)[[20, 30]].sum()), ) def test_nsmallest(self): pdf = pd.DataFrame( { "a": [1, 1, 1, 2, 2, 2, 3, 3, 3] * 3, "b": [1, 2, 2, 2, 3, 3, 3, 4, 4] * 3, "c": [1, 2, 2, 2, 3, 3, 3, 4, 4] * 3, "d": [1, 2, 2, 2, 3, 3, 3, 4, 4] * 3, }, index=np.random.rand(9 * 3), ) psdf = ps.from_pandas(pdf) self.assert_eq( psdf.groupby(["a"])["b"].nsmallest(1).sort_values(), pdf.groupby(["a"])["b"].nsmallest(1).sort_values(), ) self.assert_eq( psdf.groupby(["a"])["b"].nsmallest(2).sort_index(), pdf.groupby(["a"])["b"].nsmallest(2).sort_index(), ) self.assert_eq( (psdf.b * 10).groupby(psdf.a).nsmallest(2).sort_index(), (pdf.b * 10).groupby(pdf.a).nsmallest(2).sort_index(), ) self.assert_eq( psdf.b.rename().groupby(psdf.a).nsmallest(2).sort_index(), pdf.b.rename().groupby(pdf.a).nsmallest(2).sort_index(), ) self.assert_eq( psdf.b.groupby(psdf.a.rename()).nsmallest(2).sort_index(), pdf.b.groupby(pdf.a.rename()).nsmallest(2).sort_index(), ) self.assert_eq( psdf.b.rename().groupby(psdf.a.rename()).nsmallest(2).sort_index(), pdf.b.rename().groupby(pdf.a.rename()).nsmallest(2).sort_index(), ) with self.assertRaisesRegex(ValueError, "nsmallest do not support multi-index now"): psdf.set_index(["a", "b"]).groupby(["c"])["d"].nsmallest(1) def test_nlargest(self): pdf = pd.DataFrame( { "a": [1, 1, 1, 2, 2, 2, 3, 3, 3] * 3, "b": [1, 2, 2, 2, 3, 3, 3, 4, 4] * 3, "c": [1, 2, 2, 2, 3, 3, 3, 4, 4] * 3, "d": [1, 2, 2, 2, 3, 3, 3, 4, 4] * 3, }, index=np.random.rand(9 * 3), ) psdf = ps.from_pandas(pdf) self.assert_eq( psdf.groupby(["a"])["b"].nlargest(1).sort_values(), pdf.groupby(["a"])["b"].nlargest(1).sort_values(), ) self.assert_eq( psdf.groupby(["a"])["b"].nlargest(2).sort_index(), pdf.groupby(["a"])["b"].nlargest(2).sort_index(), ) self.assert_eq( (psdf.b * 10).groupby(psdf.a).nlargest(2).sort_index(), (pdf.b * 10).groupby(pdf.a).nlargest(2).sort_index(), ) self.assert_eq( psdf.b.rename().groupby(psdf.a).nlargest(2).sort_index(), pdf.b.rename().groupby(pdf.a).nlargest(2).sort_index(), ) self.assert_eq( psdf.b.groupby(psdf.a.rename()).nlargest(2).sort_index(), pdf.b.groupby(pdf.a.rename()).nlargest(2).sort_index(), ) self.assert_eq( psdf.b.rename().groupby(psdf.a.rename()).nlargest(2).sort_index(), pdf.b.rename().groupby(pdf.a.rename()).nlargest(2).sort_index(), ) with self.assertRaisesRegex(ValueError, "nlargest do not support multi-index now"): psdf.set_index(["a", "b"]).groupby(["c"])["d"].nlargest(1) def test_shift(self): pdf = pd.DataFrame( { "a": [1, 1, 2, 2, 3, 3] * 3, "b": [1, 1, 2, 2, 3, 4] * 3, "c": [1, 4, 9, 16, 25, 36] * 3, }, index=np.random.rand(6 * 3), ) psdf = ps.from_pandas(pdf) self.assert_eq( psdf.groupby("a").shift().sort_index(), pdf.groupby("a").shift().sort_index() ) # TODO: seems like a pandas' bug when fill_value is not None? # self.assert_eq(psdf.groupby(['a', 'b']).shift(periods=-1, fill_value=0).sort_index(), # pdf.groupby(['a', 'b']).shift(periods=-1, fill_value=0).sort_index()) self.assert_eq( psdf.groupby(["b"])["a"].shift().sort_index(), pdf.groupby(["b"])["a"].shift().sort_index(), ) self.assert_eq( psdf.groupby(["a", "b"])["c"].shift().sort_index(), pdf.groupby(["a", "b"])["c"].shift().sort_index(), ) self.assert_eq( psdf.groupby(psdf.b // 5).shift().sort_index(), pdf.groupby(pdf.b // 5).shift().sort_index(), ) self.assert_eq( psdf.groupby(psdf.b // 5)["a"].shift().sort_index(), pdf.groupby(pdf.b // 5)["a"].shift().sort_index(), ) self.assert_eq( psdf.a.rename().groupby(psdf.b).shift().sort_index(), pdf.a.rename().groupby(pdf.b).shift().sort_index(), ) self.assert_eq( psdf.a.groupby(psdf.b.rename()).shift().sort_index(), pdf.a.groupby(pdf.b.rename()).shift().sort_index(), ) self.assert_eq( psdf.a.rename().groupby(psdf.b.rename()).shift().sort_index(), pdf.a.rename().groupby(pdf.b.rename()).shift().sort_index(), ) self.assert_eq(psdf.groupby("a").shift().sum(), pdf.groupby("a").shift().sum().astype(int)) self.assert_eq( psdf.a.rename().groupby(psdf.b).shift().sum(), pdf.a.rename().groupby(pdf.b).shift().sum(), ) # 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", "a")).shift().sort_index(), pdf.groupby(("x", "a")).shift().sort_index(), ) # TODO: seems like a pandas' bug when fill_value is not None? # self.assert_eq(psdf.groupby([('x', 'a'), ('x', 'b')]).shift(periods=-1, # fill_value=0).sort_index(), # pdf.groupby([('x', 'a'), ('x', 'b')]).shift(periods=-1, # fill_value=0).sort_index()) def test_missing(self): psdf = ps.DataFrame({"a": [1, 2, 3, 4, 5, 6, 7, 8, 9]}) # DataFrameGroupBy functions missing_functions = inspect.getmembers( MissingPandasLikeDataFrameGroupBy, inspect.isfunction ) unsupported_functions = [ name for (name, type_) in missing_functions if type_.__name__ == "unsupported_function" ] for name in unsupported_functions: with self.assertRaisesRegex( PandasNotImplementedError, "method.*GroupBy.*{}.*not implemented( yet\\.|\\. .+)".format(name), ): getattr(psdf.groupby("a"), name)() deprecated_functions = [ name for (name, type_) in missing_functions if type_.__name__ == "deprecated_function" ] for name in deprecated_functions: with self.assertRaisesRegex( PandasNotImplementedError, "method.*GroupBy.*{}.*is deprecated".format(name) ): getattr(psdf.groupby("a"), name)() # SeriesGroupBy functions missing_functions = inspect.getmembers(MissingPandasLikeSeriesGroupBy, inspect.isfunction) unsupported_functions = [ name for (name, type_) in missing_functions if type_.__name__ == "unsupported_function" ] for name in unsupported_functions: with self.assertRaisesRegex( PandasNotImplementedError, "method.*GroupBy.*{}.*not implemented( yet\\.|\\. .+)".format(name), ): getattr(psdf.a.groupby(psdf.a), name)() deprecated_functions = [ name for (name, type_) in missing_functions if type_.__name__ == "deprecated_function" ] for name in deprecated_functions: with self.assertRaisesRegex( PandasNotImplementedError, "method.*GroupBy.*{}.*is deprecated".format(name) ): getattr(psdf.a.groupby(psdf.a), name)() # DataFrameGroupBy properties missing_properties = inspect.getmembers( MissingPandasLikeDataFrameGroupBy, lambda o: isinstance(o, property) ) unsupported_properties = [ name for (name, type_) in missing_properties if type_.fget.__name__ == "unsupported_property" ] for name in unsupported_properties: with self.assertRaisesRegex( PandasNotImplementedError, "property.*GroupBy.*{}.*not implemented( yet\\.|\\. .+)".format(name), ): getattr(psdf.groupby("a"), name) deprecated_properties = [ name for (name, type_) in missing_properties if type_.fget.__name__ == "deprecated_property" ] for name in deprecated_properties: with self.assertRaisesRegex( PandasNotImplementedError, "property.*GroupBy.*{}.*is deprecated".format(name) ): getattr(psdf.groupby("a"), name) # SeriesGroupBy properties missing_properties = inspect.getmembers( MissingPandasLikeSeriesGroupBy, lambda o: isinstance(o, property) ) unsupported_properties = [ name for (name, type_) in missing_properties if type_.fget.__name__ == "unsupported_property" ] for name in unsupported_properties: with self.assertRaisesRegex( PandasNotImplementedError, "property.*GroupBy.*{}.*not implemented( yet\\.|\\. .+)".format(name), ): getattr(psdf.a.groupby(psdf.a), name) deprecated_properties = [ name for (name, type_) in missing_properties if type_.fget.__name__ == "deprecated_property" ] for name in deprecated_properties: with self.assertRaisesRegex( PandasNotImplementedError, "property.*GroupBy.*{}.*is deprecated".format(name) ): getattr(psdf.a.groupby(psdf.a), name) @staticmethod def test_is_multi_agg_with_relabel(): assert is_multi_agg_with_relabel(a="max") is False assert is_multi_agg_with_relabel(a_min=("a", "max"), a_max=("a", "min")) is True def test_get_group(self): pdf = pd.DataFrame( [ ("falcon", "bird", 389.0), ("parrot", "bird", 24.0), ("lion", "mammal", 80.5), ("monkey", "mammal", np.nan), ], columns=["name", "class", "max_speed"], index=[0, 2, 3, 1], ) pdf.columns.name = "Koalas" psdf = ps.from_pandas(pdf) self.assert_eq( psdf.groupby("class").get_group("bird"), pdf.groupby("class").get_group("bird"), ) self.assert_eq( psdf.groupby("class")["name"].get_group("mammal"), pdf.groupby("class")["name"].get_group("mammal"), ) self.assert_eq( psdf.groupby("class")[["name"]].get_group("mammal"), pdf.groupby("class")[["name"]].get_group("mammal"), ) self.assert_eq( psdf.groupby(["class", "name"]).get_group(("mammal", "lion")), pdf.groupby(["class", "name"]).get_group(("mammal", "lion")), ) self.assert_eq( psdf.groupby(["class", "name"])["max_speed"].get_group(("mammal", "lion")), pdf.groupby(["class", "name"])["max_speed"].get_group(("mammal", "lion")), ) self.assert_eq( psdf.groupby(["class", "name"])[["max_speed"]].get_group(("mammal", "lion")), pdf.groupby(["class", "name"])[["max_speed"]].get_group(("mammal", "lion")), ) self.assert_eq( (psdf.max_speed + 1).groupby(psdf["class"]).get_group("mammal"), (pdf.max_speed + 1).groupby(pdf["class"]).get_group("mammal"), ) self.assert_eq( psdf.groupby("max_speed").get_group(80.5), pdf.groupby("max_speed").get_group(80.5), ) self.assertRaises(KeyError, lambda: psdf.groupby("class").get_group("fish")) self.assertRaises(TypeError, lambda: psdf.groupby("class").get_group(["bird", "mammal"])) self.assertRaises(KeyError, lambda: psdf.groupby("class")["name"].get_group("fish")) self.assertRaises( TypeError, lambda: psdf.groupby("class")["name"].get_group(["bird", "mammal"]) ) self.assertRaises( KeyError, lambda: psdf.groupby(["class", "name"]).get_group(("lion", "mammal")) ) self.assertRaises(ValueError, lambda: psdf.groupby(["class", "name"]).get_group(("lion",))) self.assertRaises( ValueError, lambda: psdf.groupby(["class", "name"]).get_group(("mammal",)) ) self.assertRaises(ValueError, lambda: psdf.groupby(["class", "name"]).get_group("mammal")) # MultiIndex columns pdf.columns = pd.MultiIndex.from_tuples([("A", "name"), ("B", "class"), ("C", "max_speed")]) pdf.columns.names = ["Hello", "Koalas"] psdf = ps.from_pandas(pdf) self.assert_eq( psdf.groupby(("B", "class")).get_group("bird"), pdf.groupby(("B", "class")).get_group("bird"), ) self.assert_eq( psdf.groupby(("B", "class"))[[("A", "name")]].get_group("mammal"), pdf.groupby(("B", "class"))[[("A", "name")]].get_group("mammal"), ) self.assert_eq( psdf.groupby([("B", "class"), ("A", "name")]).get_group(("mammal", "lion")), pdf.groupby([("B", "class"), ("A", "name")]).get_group(("mammal", "lion")), ) self.assert_eq( psdf.groupby([("B", "class"), ("A", "name")])[[("C", "max_speed")]].get_group( ("mammal", "lion") ), pdf.groupby([("B", "class"), ("A", "name")])[[("C", "max_speed")]].get_group( ("mammal", "lion") ), ) self.assert_eq( (psdf[("C", "max_speed")] + 1).groupby(psdf[("B", "class")]).get_group("mammal"), (pdf[("C", "max_speed")] + 1).groupby(pdf[("B", "class")]).get_group("mammal"), ) self.assert_eq( psdf.groupby(("C", "max_speed")).get_group(80.5), pdf.groupby(("C", "max_speed")).get_group(80.5), ) self.assertRaises(KeyError, lambda: psdf.groupby(("B", "class")).get_group("fish")) self.assertRaises( TypeError, lambda: psdf.groupby(("B", "class")).get_group(["bird", "mammal"]) ) self.assertRaises( KeyError, lambda: psdf.groupby(("B", "class"))[("A", "name")].get_group("fish") ) self.assertRaises( KeyError, lambda: psdf.groupby([("B", "class"), ("A", "name")]).get_group(("lion", "mammal")), ) self.assertRaises( ValueError, lambda: psdf.groupby([("B", "class"), ("A", "name")]).get_group(("lion",)), ) self.assertRaises( ValueError, lambda: psdf.groupby([("B", "class"), ("A", "name")]).get_group(("mammal",)) ) self.assertRaises( ValueError, lambda: psdf.groupby([("B", "class"), ("A", "name")]).get_group("mammal") ) def test_getitem(self): psdf = ps.DataFrame( { "a": [1, 1, 1, 1, 2, 2, 2, 3, 3, 3] * 3, "b": [2, 3, 1, 4, 6, 9, 8, 10, 7, 5] * 3, "c": [3, 5, 2, 5, 1, 2, 6, 4, 3, 6] * 3, }, index=np.random.rand(10 * 3), ) self.assertTrue(isinstance(psdf.groupby("a")["b"], SeriesGroupBy)) def test_all_any(self): pdf = pd.DataFrame( { "A": [1, 1, 2, 2, 3, 3, 4, 4, 5, 5], "B": [True, True, True, False, False, False, None, True, None, False], } ) psdf = ps.from_pandas(pdf) for as_index in [True, False]: if as_index: def sort(df): return df.sort_index() else: def sort(df): return df.sort_values("A").reset_index(drop=True) self.assert_eq( sort(psdf.groupby("A", as_index=as_index).all()), sort(pdf.groupby("A", as_index=as_index).all()), ) self.assert_eq( sort(psdf.groupby("A", as_index=as_index).any()), sort(pdf.groupby("A", as_index=as_index).any()), ) self.assert_eq( sort(psdf.groupby("A", as_index=as_index).all()).B, sort(pdf.groupby("A", as_index=as_index).all()).B, ) self.assert_eq( sort(psdf.groupby("A", as_index=as_index).any()).B, sort(pdf.groupby("A", as_index=as_index).any()).B, ) self.assert_eq( psdf.B.groupby(psdf.A).all().sort_index(), pdf.B.groupby(pdf.A).all().sort_index() ) self.assert_eq( psdf.B.groupby(psdf.A).any().sort_index(), pdf.B.groupby(pdf.A).any().sort_index() ) # multi-index columns columns = pd.MultiIndex.from_tuples([("X", "A"), ("Y", "B")]) pdf.columns = columns psdf.columns = columns for as_index in [True, False]: if as_index: def sort(df): return df.sort_index() else: def sort(df): return df.sort_values(("X", "A")).reset_index(drop=True) self.assert_eq( sort(psdf.groupby(("X", "A"), as_index=as_index).all()), sort(pdf.groupby(("X", "A"), as_index=as_index).all()), ) self.assert_eq( sort(psdf.groupby(("X", "A"), as_index=as_index).any()), sort(pdf.groupby(("X", "A"), as_index=as_index).any()), ) # Test skipna pdf = pd.DataFrame({"A": [True, True], "B": [1, np.nan], "C": [True, None]}) pdf.name = "x" psdf = ps.from_pandas(pdf) self.assert_eq( psdf.groupby("A").all(skipna=False).sort_index(), pdf.groupby("A").all(skipna=False).sort_index(), ) self.assert_eq( psdf.groupby("A").all(skipna=True).sort_index(), pdf.groupby("A").all(skipna=True).sort_index(), ) def test_raises(self): psdf = ps.DataFrame( {"a": [1, 2, 6, 4, 4, 6, 4, 3, 7], "b": [4, 2, 7, 3, 3, 1, 1, 1, 2]}, index=[0, 1, 3, 5, 6, 8, 9, 9, 9], ) # test raises with incorrect key self.assertRaises(ValueError, lambda: psdf.groupby([])) self.assertRaises(KeyError, lambda: psdf.groupby("x")) self.assertRaises(KeyError, lambda: psdf.groupby(["a", "x"])) self.assertRaises(KeyError, lambda: psdf.groupby("a")["x"]) self.assertRaises(KeyError, lambda: psdf.groupby("a")["b", "x"]) self.assertRaises(KeyError, lambda: psdf.groupby("a")[["b", "x"]]) def test_nunique(self): pdf = pd.DataFrame( {"a": [1, 1, 1, 1, 1, 0, 0, 0, 0, 0], "b": [2, 2, 2, 3, 3, 4, 4, 5, 5, 5]} ) psdf = ps.from_pandas(pdf) self.assert_eq( psdf.groupby("a").agg({"b": "nunique"}).sort_index(), pdf.groupby("a").agg({"b": "nunique"}).sort_index(), ) if LooseVersion(pd.__version__) < LooseVersion("1.1.0"): expected = ps.DataFrame({"b": [2, 2]}, index=pd.Index([0, 1], name="a")) self.assert_eq(psdf.groupby("a").nunique().sort_index(), expected) self.assert_eq( psdf.groupby("a").nunique(dropna=False).sort_index(), expected, ) else: self.assert_eq( psdf.groupby("a").nunique().sort_index(), pdf.groupby("a").nunique().sort_index() ) self.assert_eq( psdf.groupby("a").nunique(dropna=False).sort_index(), pdf.groupby("a").nunique(dropna=False).sort_index(), ) self.assert_eq( psdf.groupby("a")["b"].nunique().sort_index(), pdf.groupby("a")["b"].nunique().sort_index(), ) self.assert_eq( psdf.groupby("a")["b"].nunique(dropna=False).sort_index(), pdf.groupby("a")["b"].nunique(dropna=False).sort_index(), ) nunique_psdf = psdf.groupby("a", as_index=False).agg({"b": "nunique"}) nunique_pdf = pdf.groupby("a", as_index=False).agg({"b": "nunique"}) self.assert_eq( nunique_psdf.sort_values(["a", "b"]).reset_index(drop=True), nunique_pdf.sort_values(["a", "b"]).reset_index(drop=True), ) # multi-index columns columns = pd.MultiIndex.from_tuples([("x", "a"), ("y", "b")]) pdf.columns = columns psdf.columns = columns if LooseVersion(pd.__version__) < LooseVersion("1.1.0"): expected = ps.DataFrame({("y", "b"): [2, 2]}, index=pd.Index([0, 1], name=("x", "a"))) self.assert_eq( psdf.groupby(("x", "a")).nunique().sort_index(), expected, ) self.assert_eq( psdf.groupby(("x", "a")).nunique(dropna=False).sort_index(), expected, ) else: self.assert_eq( psdf.groupby(("x", "a")).nunique().sort_index(), pdf.groupby(("x", "a")).nunique().sort_index(), ) self.assert_eq( psdf.groupby(("x", "a")).nunique(dropna=False).sort_index(), pdf.groupby(("x", "a")).nunique(dropna=False).sort_index(), ) def test_unique(self): for pdf in [ pd.DataFrame( {"a": [1, 1, 1, 1, 1, 0, 0, 0, 0, 0], "b": [2, 2, 2, 3, 3, 4, 4, 5, 5, 5]} ), pd.DataFrame( { "a": [1, 1, 1, 1, 1, 0, 0, 0, 0, 0], "b": ["w", "w", "w", "x", "x", "y", "y", "z", "z", "z"], } ), ]: with self.subTest(pdf=pdf): psdf = ps.from_pandas(pdf) actual = psdf.groupby("a")["b"].unique().sort_index()._to_pandas() expect = pdf.groupby("a")["b"].unique().sort_index() self.assert_eq(len(actual), len(expect)) for act, exp in zip(actual, expect): self.assertTrue(sorted(act) == sorted(exp)) @unittest.skipIf( LooseVersion(pd.__version__) >= LooseVersion("2.0.0"), "TODO(SPARK-43444): Enable GroupBySlowTests.test_value_counts for pandas 2.0.0.", ) def test_value_counts(self): pdf = pd.DataFrame( {"A": [np.nan, 2, 2, 3, 3, 3], "B": [1, 1, 2, 3, 3, np.nan]}, columns=["A", "B"] ) psdf = ps.from_pandas(pdf) self.assert_eq( psdf.groupby("A")["B"].value_counts().sort_index(), pdf.groupby("A")["B"].value_counts().sort_index(), ) self.assert_eq( psdf.groupby("A")["B"].value_counts(dropna=False).sort_index(), pdf.groupby("A")["B"].value_counts(dropna=False).sort_index(), ) self.assert_eq( psdf.groupby("A", dropna=False)["B"].value_counts(dropna=False).sort_index(), pdf.groupby("A", dropna=False)["B"].value_counts(dropna=False).sort_index(), # Returns are the same considering values and types, # disable check_exact to pass the assert_eq check_exact=False, ) self.assert_eq( psdf.groupby("A")["B"].value_counts(sort=True, ascending=False).sort_index(), pdf.groupby("A")["B"].value_counts(sort=True, ascending=False).sort_index(), ) self.assert_eq( psdf.groupby("A")["B"] .value_counts(sort=True, ascending=False, dropna=False) .sort_index(), pdf.groupby("A")["B"] .value_counts(sort=True, ascending=False, dropna=False) .sort_index(), ) self.assert_eq( psdf.groupby("A")["B"] .value_counts(sort=True, ascending=True, dropna=False) .sort_index(), pdf.groupby("A")["B"] .value_counts(sort=True, ascending=True, dropna=False) .sort_index(), ) self.assert_eq( psdf.B.rename().groupby(psdf.A).value_counts().sort_index(), pdf.B.rename().groupby(pdf.A).value_counts().sort_index(), ) self.assert_eq( psdf.B.rename().groupby(psdf.A, dropna=False).value_counts().sort_index(), pdf.B.rename().groupby(pdf.A, dropna=False).value_counts().sort_index(), # Returns are the same considering values and types, # disable check_exact to pass the assert_eq check_exact=False, ) self.assert_eq( psdf.B.groupby(psdf.A.rename()).value_counts().sort_index(), pdf.B.groupby(pdf.A.rename()).value_counts().sort_index(), ) self.assert_eq( psdf.B.rename().groupby(psdf.A.rename()).value_counts().sort_index(), pdf.B.rename().groupby(pdf.A.rename()).value_counts().sort_index(), ) def test_size(self): pdf = pd.DataFrame({"A": [1, 2, 2, 3, 3, 3], "B": [1, 1, 2, 3, 3, 3]}) psdf = ps.from_pandas(pdf) self.assert_eq(psdf.groupby("A").size().sort_index(), pdf.groupby("A").size().sort_index()) self.assert_eq( psdf.groupby("A")["B"].size().sort_index(), pdf.groupby("A")["B"].size().sort_index() ) self.assert_eq( psdf.groupby("A")[["B"]].size().sort_index(), pdf.groupby("A")[["B"]].size().sort_index(), ) self.assert_eq( psdf.groupby(["A", "B"]).size().sort_index(), pdf.groupby(["A", "B"]).size().sort_index(), ) # multi-index columns columns = pd.MultiIndex.from_tuples([("X", "A"), ("Y", "B")]) pdf.columns = columns psdf.columns = columns self.assert_eq( psdf.groupby(("X", "A")).size().sort_index(), pdf.groupby(("X", "A")).size().sort_index(), ) self.assert_eq( psdf.groupby([("X", "A"), ("Y", "B")]).size().sort_index(), pdf.groupby([("X", "A"), ("Y", "B")]).size().sort_index(), ) def test_diff(self): pdf = pd.DataFrame( { "a": [1, 2, 3, 4, 5, 6] * 3, "b": [1, 1, 2, 3, 5, 8] * 3, "c": [1, 4, 9, 16, 25, 36] * 3, } ) psdf = ps.from_pandas(pdf) self.assert_eq(psdf.groupby("b").diff().sort_index(), pdf.groupby("b").diff().sort_index()) self.assert_eq( psdf.groupby(["a", "b"]).diff().sort_index(), pdf.groupby(["a", "b"]).diff().sort_index(), ) self.assert_eq( psdf.groupby(["b"])["a"].diff().sort_index(), pdf.groupby(["b"])["a"].diff().sort_index(), ) self.assert_eq( psdf.groupby(["b"])[["a", "b"]].diff().sort_index(), pdf.groupby(["b"])[["a", "b"]].diff().sort_index(), ) self.assert_eq( psdf.groupby(psdf.b // 5).diff().sort_index(), pdf.groupby(pdf.b // 5).diff().sort_index(), ) self.assert_eq( psdf.groupby(psdf.b // 5)["a"].diff().sort_index(), pdf.groupby(pdf.b // 5)["a"].diff().sort_index(), ) self.assert_eq(psdf.groupby("b").diff().sum(), pdf.groupby("b").diff().sum().astype(int)) self.assert_eq(psdf.groupby(["b"])["a"].diff().sum(), pdf.groupby(["b"])["a"].diff().sum()) # 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")).diff().sort_index(), pdf.groupby(("x", "b")).diff().sort_index(), ) self.assert_eq( psdf.groupby([("x", "a"), ("x", "b")]).diff().sort_index(), pdf.groupby([("x", "a"), ("x", "b")]).diff().sort_index(), ) def test_rank(self): pdf = pd.DataFrame( { "a": [1, 2, 3, 4, 5, 6] * 3, "b": [1, 1, 2, 3, 5, 8] * 3, "c": [1, 4, 9, 16, 25, 36] * 3, }, index=np.random.rand(6 * 3), ) psdf = ps.from_pandas(pdf) self.assert_eq(psdf.groupby("b").rank().sort_index(), pdf.groupby("b").rank().sort_index()) self.assert_eq( psdf.groupby(["a", "b"]).rank().sort_index(), pdf.groupby(["a", "b"]).rank().sort_index(), ) self.assert_eq( psdf.groupby(["b"])["a"].rank().sort_index(), pdf.groupby(["b"])["a"].rank().sort_index(), ) self.assert_eq( psdf.groupby(["b"])[["a", "c"]].rank().sort_index(), pdf.groupby(["b"])[["a", "c"]].rank().sort_index(), ) self.assert_eq( psdf.groupby(psdf.b // 5).rank().sort_index(), pdf.groupby(pdf.b // 5).rank().sort_index(), ) self.assert_eq( psdf.groupby(psdf.b // 5)["a"].rank().sort_index(), pdf.groupby(pdf.b // 5)["a"].rank().sort_index(), ) self.assert_eq(psdf.groupby("b").rank().sum(), pdf.groupby("b").rank().sum()) self.assert_eq(psdf.groupby(["b"])["a"].rank().sum(), pdf.groupby(["b"])["a"].rank().sum()) # 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")).rank().sort_index(), pdf.groupby(("x", "b")).rank().sort_index(), ) self.assert_eq( psdf.groupby([("x", "a"), ("x", "b")]).rank().sort_index(), pdf.groupby([("x", "a"), ("x", "b")]).rank().sort_index(), ) class GroupByTests(GroupByTestsMixin, PandasOnSparkTestCase, TestUtils): pass if __name__ == "__main__": from pyspark.pandas.tests.groupby.test_groupby import * # noqa: F401 try: import xmlrunner testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) except ImportError: testRunner = None unittest.main(testRunner=testRunner, verbosity=2)