| # |
| # 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 numpy as np |
| import pandas as pd |
| |
| from pyspark import pandas as ps |
| from pyspark.testing.pandasutils import PandasOnSparkTestCase |
| from pyspark.testing.sqlutils import SQLTestUtils |
| |
| |
| class IndexingLocMixin: |
| @property |
| def pdf(self): |
| return pd.DataFrame( |
| {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]}, |
| index=[0, 1, 3, 5, 6, 8, 9, 9, 9], |
| ) |
| |
| @property |
| def pdf2(self): |
| return pd.DataFrame( |
| {0: [1, 2, 3, 4, 5, 6, 7, 8, 9], 1: [4, 5, 6, 3, 2, 1, 0, 0, 0]}, |
| index=[0, 1, 3, 5, 6, 8, 9, 9, 9], |
| ) |
| |
| @property |
| def psdf(self): |
| return ps.from_pandas(self.pdf) |
| |
| @property |
| def psdf2(self): |
| return ps.from_pandas(self.pdf2) |
| |
| def test_loc(self): |
| psdf = self.psdf |
| pdf = self.pdf |
| |
| self.assert_eq(psdf.loc[5:5], pdf.loc[5:5]) |
| self.assert_eq(psdf.loc[3:8], pdf.loc[3:8]) |
| self.assert_eq(psdf.loc[:8], pdf.loc[:8]) |
| self.assert_eq(psdf.loc[3:], pdf.loc[3:]) |
| self.assert_eq(psdf.loc[[5]], pdf.loc[[5]]) |
| self.assert_eq(psdf.loc[:], pdf.loc[:]) |
| |
| # TODO?: self.assert_eq(psdf.loc[[3, 4, 1, 8]], pdf.loc[[3, 4, 1, 8]]) |
| # TODO?: self.assert_eq(psdf.loc[[3, 4, 1, 9]], pdf.loc[[3, 4, 1, 9]]) |
| # TODO?: self.assert_eq(psdf.loc[np.array([3, 4, 1, 9])], pdf.loc[np.array([3, 4, 1, 9])]) |
| |
| self.assert_eq(psdf.a.loc[5:5], pdf.a.loc[5:5]) |
| self.assert_eq(psdf.a.loc[3:8], pdf.a.loc[3:8]) |
| self.assert_eq(psdf.a.loc[:8], pdf.a.loc[:8]) |
| self.assert_eq(psdf.a.loc[3:], pdf.a.loc[3:]) |
| self.assert_eq(psdf.a.loc[[5]], pdf.a.loc[[5]]) |
| |
| # TODO?: self.assert_eq(psdf.a.loc[[3, 4, 1, 8]], pdf.a.loc[[3, 4, 1, 8]]) |
| # TODO?: self.assert_eq(psdf.a.loc[[3, 4, 1, 9]], pdf.a.loc[[3, 4, 1, 9]]) |
| # TODO?: self.assert_eq(psdf.a.loc[np.array([3, 4, 1, 9])], |
| # pdf.a.loc[np.array([3, 4, 1, 9])]) |
| |
| self.assert_eq(psdf.a.loc[[]], pdf.a.loc[[]]) |
| self.assert_eq(psdf.a.loc[np.array([])], pdf.a.loc[np.array([])]) |
| |
| self.assert_eq(psdf.loc[1000:], pdf.loc[1000:]) |
| self.assert_eq(psdf.loc[-2000:-1000], pdf.loc[-2000:-1000]) |
| |
| self.assert_eq(psdf.loc[5], pdf.loc[5]) |
| self.assert_eq(psdf.loc[9], pdf.loc[9]) |
| self.assert_eq(psdf.a.loc[5], pdf.a.loc[5]) |
| self.assert_eq(psdf.a.loc[9], pdf.a.loc[9]) |
| |
| self.assertRaises(KeyError, lambda: psdf.loc[10]) |
| self.assertRaises(KeyError, lambda: psdf.a.loc[10]) |
| |
| # monotonically increasing index test |
| pdf = pd.DataFrame({"a": [1, 2, 3, 4, 5, 6, 7, 8, 9]}, index=[0, 1, 1, 2, 2, 2, 4, 5, 6]) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(psdf.loc[:2], pdf.loc[:2]) |
| self.assert_eq(psdf.loc[:3], pdf.loc[:3]) |
| self.assert_eq(psdf.loc[3:], pdf.loc[3:]) |
| self.assert_eq(psdf.loc[4:], pdf.loc[4:]) |
| self.assert_eq(psdf.loc[3:2], pdf.loc[3:2]) |
| self.assert_eq(psdf.loc[-1:2], pdf.loc[-1:2]) |
| self.assert_eq(psdf.loc[3:10], pdf.loc[3:10]) |
| |
| # monotonically decreasing index test |
| pdf = pd.DataFrame({"a": [1, 2, 3, 4, 5, 6, 7, 8, 9]}, index=[6, 5, 5, 4, 4, 4, 2, 1, 0]) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(psdf.loc[:4], pdf.loc[:4]) |
| self.assert_eq(psdf.loc[:3], pdf.loc[:3]) |
| self.assert_eq(psdf.loc[3:], pdf.loc[3:]) |
| self.assert_eq(psdf.loc[2:], pdf.loc[2:]) |
| self.assert_eq(psdf.loc[2:3], pdf.loc[2:3]) |
| self.assert_eq(psdf.loc[2:-1], pdf.loc[2:-1]) |
| self.assert_eq(psdf.loc[10:3], pdf.loc[10:3]) |
| |
| # test when type of key is string and given value is not included in key |
| pdf = pd.DataFrame({"a": [1, 2, 3]}, index=["a", "b", "d"]) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(psdf.loc["a":"z"], pdf.loc["a":"z"]) |
| |
| # KeyError when index is not monotonic increasing or decreasing |
| # and specified values don't exist in index |
| psdf = ps.DataFrame([[1, 2], [4, 5], [7, 8]], index=["cobra", "viper", "sidewinder"]) |
| |
| self.assertRaises(KeyError, lambda: psdf.loc["cobra":"koalas"]) |
| self.assertRaises(KeyError, lambda: psdf.loc["koalas":"viper"]) |
| |
| psdf = ps.DataFrame([[1, 2], [4, 5], [7, 8]], index=[10, 30, 20]) |
| |
| self.assertRaises(KeyError, lambda: psdf.loc[0:30]) |
| self.assertRaises(KeyError, lambda: psdf.loc[10:100]) |
| |
| def test_loc_getitem_boolean_series(self): |
| pdf = pd.DataFrame( |
| {"A": [0, 1, 2, 3, 4], "B": [100, 200, 300, 400, 500]}, index=[20, 10, 30, 0, 50] |
| ) |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq(pdf.A.loc[pdf.B > 200], psdf.A.loc[psdf.B > 200]) |
| self.assert_eq(pdf.B.loc[pdf.B > 200], psdf.B.loc[psdf.B > 200]) |
| self.assert_eq(pdf.loc[pdf.B > 200], psdf.loc[psdf.B > 200]) |
| |
| def test_loc_non_informative_index(self): |
| pdf = pd.DataFrame({"x": [1, 2, 3, 4]}, index=[10, 20, 30, 40]) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(psdf.loc[20:30], pdf.loc[20:30]) |
| |
| pdf = pd.DataFrame({"x": [1, 2, 3, 4]}, index=[10, 20, 20, 40]) |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq(psdf.loc[20:20], pdf.loc[20:20]) |
| |
| def test_loc_with_series(self): |
| psdf = self.psdf |
| pdf = self.pdf |
| |
| self.assert_eq(psdf.loc[psdf.a % 2 == 0], pdf.loc[pdf.a % 2 == 0]) |
| self.assert_eq(psdf.loc[psdf.a % 2 == 0, "a"], pdf.loc[pdf.a % 2 == 0, "a"]) |
| self.assert_eq(psdf.loc[psdf.a % 2 == 0, ["a"]], pdf.loc[pdf.a % 2 == 0, ["a"]]) |
| self.assert_eq(psdf.a.loc[psdf.a % 2 == 0], pdf.a.loc[pdf.a % 2 == 0]) |
| |
| self.assert_eq(psdf.loc[psdf.copy().a % 2 == 0], pdf.loc[pdf.copy().a % 2 == 0]) |
| self.assert_eq(psdf.loc[psdf.copy().a % 2 == 0, "a"], pdf.loc[pdf.copy().a % 2 == 0, "a"]) |
| self.assert_eq( |
| psdf.loc[psdf.copy().a % 2 == 0, ["a"]], pdf.loc[pdf.copy().a % 2 == 0, ["a"]] |
| ) |
| self.assert_eq(psdf.a.loc[psdf.copy().a % 2 == 0], pdf.a.loc[pdf.copy().a % 2 == 0]) |
| |
| def test_loc_noindex(self): |
| psdf = self.psdf |
| psdf = psdf.reset_index() |
| pdf = self.pdf |
| pdf = pdf.reset_index() |
| |
| self.assert_eq(psdf[["a"]], pdf[["a"]]) |
| |
| self.assert_eq(psdf.loc[:], pdf.loc[:]) |
| self.assert_eq(psdf.loc[5:5], pdf.loc[5:5]) |
| |
| def test_loc_on_numpy_datetimes(self): |
| pdf = pd.DataFrame( |
| {"x": [1, 2, 3]}, index=list(map(np.datetime64, ["2014", "2015", "2016"])) |
| ) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(psdf.loc["2014":"2015"], pdf.loc["2014":"2015"]) |
| |
| def test_loc_on_pandas_datetimes(self): |
| pdf = pd.DataFrame( |
| {"x": [1, 2, 3]}, index=list(map(pd.Timestamp, ["2014", "2015", "2016"])) |
| ) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(psdf.loc["2014":"2015"], pdf.loc["2014":"2015"]) |
| |
| @unittest.skip("TODO?: the behavior of slice for datetime") |
| def test_loc_datetime_no_freq(self): |
| datetime_index = pd.date_range("2016-01-01", "2016-01-31", freq="12h") |
| datetime_index.freq = None # FORGET FREQUENCY |
| pdf = pd.DataFrame({"num": range(len(datetime_index))}, index=datetime_index) |
| psdf = ps.from_pandas(pdf) |
| |
| slice_ = slice("2016-01-03", "2016-01-05") |
| result = psdf.loc[slice_, :] |
| expected = pdf.loc[slice_, :] |
| self.assert_eq(result, expected) |
| |
| @unittest.skip("TODO?: the behavior of slice for datetime") |
| def test_loc_timestamp_str(self): |
| pdf = pd.DataFrame( |
| {"A": np.random.randn(100), "B": np.random.randn(100)}, |
| index=pd.date_range("2011-01-01", freq="H", periods=100), |
| ) |
| psdf = ps.from_pandas(pdf) |
| |
| # partial string slice |
| # TODO?: self.assert_eq(pdf.loc['2011-01-02'], |
| # TODO?: psdf.loc['2011-01-02']) |
| self.assert_eq(pdf.loc["2011-01-02":"2011-01-05"], psdf.loc["2011-01-02":"2011-01-05"]) |
| |
| # series |
| # TODO?: self.assert_eq(pdf.A.loc['2011-01-02'], |
| # TODO?: psdf.A.loc['2011-01-02']) |
| self.assert_eq(pdf.A.loc["2011-01-02":"2011-01-05"], psdf.A.loc["2011-01-02":"2011-01-05"]) |
| |
| pdf = pd.DataFrame( |
| {"A": np.random.randn(100), "B": np.random.randn(100)}, |
| index=pd.date_range("2011-01-01", freq="M", periods=100), |
| ) |
| psdf = ps.from_pandas(pdf) |
| # TODO?: self.assert_eq(pdf.loc['2011-01'], psdf.loc['2011-01']) |
| # TODO?: self.assert_eq(pdf.loc['2011'], psdf.loc['2011']) |
| |
| self.assert_eq(pdf.loc["2011-01":"2012-05"], psdf.loc["2011-01":"2012-05"]) |
| self.assert_eq(pdf.loc["2011":"2015"], psdf.loc["2011":"2015"]) |
| |
| # series |
| # TODO?: self.assert_eq(pdf.B.loc['2011-01'], psdf.B.loc['2011-01']) |
| # TODO?: self.assert_eq(pdf.B.loc['2011'], psdf.B.loc['2011']) |
| |
| self.assert_eq(pdf.B.loc["2011-01":"2012-05"], psdf.B.loc["2011-01":"2012-05"]) |
| self.assert_eq(pdf.B.loc["2011":"2015"], psdf.B.loc["2011":"2015"]) |
| |
| def test_frame_loc_setitem(self): |
| pdf = pd.DataFrame( |
| [[1, 2], [4, 5], [7, 8]], |
| index=["cobra", "viper", "sidewinder"], |
| columns=["max_speed", "shield"], |
| ) |
| psdf = ps.from_pandas(pdf) |
| |
| pser1 = pdf.max_speed |
| pser2 = pdf.shield |
| psser1 = psdf.max_speed |
| psser2 = psdf.shield |
| |
| pdf.loc[["viper", "sidewinder"], ["shield", "max_speed"]] = 10 |
| psdf.loc[["viper", "sidewinder"], ["shield", "max_speed"]] = 10 |
| self.assert_eq(psdf, pdf) |
| self.assert_eq(psser1, pser1) |
| self.assert_eq(psser2, pser2) |
| |
| pdf.loc[["viper", "sidewinder"], "shield"] = 50 |
| psdf.loc[["viper", "sidewinder"], "shield"] = 50 |
| self.assert_eq(psdf, pdf) |
| self.assert_eq(psser1, pser1) |
| self.assert_eq(psser2, pser2) |
| |
| pdf.loc["cobra", "max_speed"] = 30 |
| psdf.loc["cobra", "max_speed"] = 30 |
| self.assert_eq(psdf, pdf) |
| self.assert_eq(psser1, pser1) |
| self.assert_eq(psser2, pser2) |
| |
| pdf.loc[pdf.max_speed < 5, "max_speed"] = -pdf.max_speed |
| psdf.loc[psdf.max_speed < 5, "max_speed"] = -psdf.max_speed |
| self.assert_eq(psdf, pdf) |
| self.assert_eq(psser1, pser1) |
| self.assert_eq(psser2, pser2) |
| |
| pdf.loc[pdf.max_speed < 2, "max_speed"] = -pdf.max_speed |
| psdf.loc[psdf.max_speed < 2, "max_speed"] = -psdf.max_speed |
| self.assert_eq(psdf, pdf) |
| self.assert_eq(psser1, pser1) |
| self.assert_eq(psser2, pser2) |
| |
| pdf.loc[:, "min_speed"] = 0 |
| psdf.loc[:, "min_speed"] = 0 |
| self.assert_eq(psdf, pdf, almost=True) |
| self.assert_eq(psser1, pser1) |
| self.assert_eq(psser2, pser2) |
| |
| with self.assertRaisesRegex(ValueError, "Incompatible indexer with Series"): |
| psdf.loc["cobra", "max_speed"] = -psdf.max_speed |
| with self.assertRaisesRegex(ValueError, "shape mismatch"): |
| psdf.loc[:, ["shield", "max_speed"]] = -psdf.max_speed |
| with self.assertRaisesRegex(ValueError, "Only a dataframe with one column can be assigned"): |
| psdf.loc[:, "max_speed"] = psdf |
| |
| # multi-index columns |
| columns = pd.MultiIndex.from_tuples( |
| [("x", "max_speed"), ("x", "shield"), ("y", "min_speed")] |
| ) |
| pdf.columns = columns |
| psdf.columns = columns |
| |
| pdf.loc[:, ("y", "shield")] = -pdf[("x", "shield")] |
| psdf.loc[:, ("y", "shield")] = -psdf[("x", "shield")] |
| self.assert_eq(psdf, pdf, almost=True) |
| self.assert_eq(psser1, pser1) |
| self.assert_eq(psser2, pser2) |
| |
| pdf.loc[:, "z"] = 100 |
| psdf.loc[:, "z"] = 100 |
| self.assert_eq(psdf, pdf, almost=True) |
| self.assert_eq(psser1, pser1) |
| self.assert_eq(psser2, pser2) |
| |
| with self.assertRaisesRegex(KeyError, "Key length \\(3\\) exceeds index depth \\(2\\)"): |
| psdf.loc[:, [("x", "max_speed", "foo")]] = -psdf[("x", "shield")] |
| |
| pdf = pd.DataFrame( |
| [[1], [4], [7]], index=["cobra", "viper", "sidewinder"], columns=["max_speed"] |
| ) |
| psdf = ps.from_pandas(pdf) |
| |
| pdf.loc[:, "max_speed"] = pdf |
| psdf.loc[:, "max_speed"] = psdf |
| self.assert_eq(psdf, pdf) |
| |
| def test_series_loc_setitem(self): |
| pdf = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]}, index=["cobra", "viper", "sidewinder"]) |
| psdf = ps.from_pandas(pdf) |
| |
| pser = pdf.x |
| psery = pdf.y |
| psser = psdf.x |
| pssery = psdf.y |
| |
| pser.loc[pser % 2 == 1] = -pser |
| psser.loc[psser % 2 == 1] = -psser |
| self.assert_eq(psser, pser) |
| self.assert_eq(psdf, pdf) |
| self.assert_eq(pssery, psery) |
| |
| for key, value in [ |
| (["viper", "sidewinder"], 10), |
| ("viper", 50), |
| (slice(None), 10), |
| (slice(None, "viper"), 20), |
| (slice("viper", None), 30), |
| ]: |
| with self.subTest(key=key, value=value): |
| pser.loc[key] = value |
| psser.loc[key] = value |
| self.assert_eq(psser, pser) |
| self.assert_eq(psdf, pdf) |
| self.assert_eq(pssery, psery) |
| |
| with self.assertRaises(ValueError): |
| psser.loc["viper"] = -psser |
| |
| # multiindex |
| pser = pd.Series( |
| [1, 2, 3], |
| index=pd.MultiIndex.from_tuples([("x", "cobra"), ("x", "viper"), ("y", "sidewinder")]), |
| ) |
| psser = ps.from_pandas(pser) |
| |
| pser.loc["x"] = pser * 10 |
| psser.loc["x"] = psser * 10 |
| self.assert_eq(psser, pser) |
| |
| pser.loc["y"] = pser * 10 |
| psser.loc["y"] = psser * 10 |
| self.assert_eq(psser, pser) |
| |
| |
| class IndexingLocTests( |
| IndexingLocMixin, |
| PandasOnSparkTestCase, |
| SQLTestUtils, |
| ): |
| pass |
| |
| |
| if __name__ == "__main__": |
| from pyspark.pandas.tests.indexes.test_indexing_loc 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) |
