| # |
| # 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 FrameRenameMixin: |
| def test_rename_dataframe(self): |
| pdf1 = pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]}) |
| psdf1 = ps.from_pandas(pdf1) |
| |
| self.assert_eq( |
| psdf1.rename(columns={"A": "a", "B": "b"}), pdf1.rename(columns={"A": "a", "B": "b"}) |
| ) |
| |
| result_psdf = psdf1.rename(index={1: 10, 2: 20}) |
| result_pdf = pdf1.rename(index={1: 10, 2: 20}) |
| self.assert_eq(result_psdf, result_pdf) |
| |
| # inplace |
| pser = result_pdf.A |
| psser = result_psdf.A |
| result_psdf.rename(index={10: 100, 20: 200}, inplace=True) |
| result_pdf.rename(index={10: 100, 20: 200}, inplace=True) |
| self.assert_eq(result_psdf, result_pdf) |
| self.assert_eq(psser, pser) |
| |
| def str_lower(s) -> str: |
| return str.lower(s) |
| |
| self.assert_eq( |
| psdf1.rename(str_lower, axis="columns"), pdf1.rename(str_lower, axis="columns") |
| ) |
| |
| def mul10(x) -> int: |
| return x * 10 |
| |
| self.assert_eq(psdf1.rename(mul10, axis="index"), pdf1.rename(mul10, axis="index")) |
| |
| self.assert_eq( |
| psdf1.rename(columns=str_lower, index={1: 10, 2: 20}), |
| pdf1.rename(columns=str_lower, index={1: 10, 2: 20}), |
| ) |
| |
| self.assert_eq( |
| psdf1.rename(columns=lambda x: str.lower(x)), |
| pdf1.rename(columns=lambda x: str.lower(x)), |
| ) |
| |
| idx = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B"), ("Y", "C"), ("Y", "D")]) |
| pdf2 = pd.DataFrame([[1, 2, 3, 4], [5, 6, 7, 8]], columns=idx) |
| psdf2 = ps.from_pandas(pdf2) |
| |
| self.assert_eq(psdf2.rename(columns=str_lower), pdf2.rename(columns=str_lower)) |
| self.assert_eq( |
| psdf2.rename(columns=lambda x: str.lower(x)), |
| pdf2.rename(columns=lambda x: str.lower(x)), |
| ) |
| |
| self.assert_eq( |
| psdf2.rename(columns=str_lower, level=0), pdf2.rename(columns=str_lower, level=0) |
| ) |
| self.assert_eq( |
| psdf2.rename(columns=str_lower, level=1), pdf2.rename(columns=str_lower, level=1) |
| ) |
| |
| pdf3 = pd.DataFrame([[1, 2], [3, 4], [5, 6], [7, 8]], index=idx, columns=list("ab")) |
| psdf3 = ps.from_pandas(pdf3) |
| |
| self.assert_eq(psdf3.rename(index=str_lower), pdf3.rename(index=str_lower)) |
| self.assert_eq( |
| psdf3.rename(index=str_lower, level=0), pdf3.rename(index=str_lower, level=0) |
| ) |
| self.assert_eq( |
| psdf3.rename(index=str_lower, level=1), pdf3.rename(index=str_lower, level=1) |
| ) |
| |
| pdf4 = pdf2 + 1 |
| psdf4 = psdf2 + 1 |
| self.assert_eq(psdf4.rename(columns=str_lower), pdf4.rename(columns=str_lower)) |
| |
| pdf5 = pdf3 + 1 |
| psdf5 = psdf3 + 1 |
| self.assert_eq(psdf5.rename(index=str_lower), pdf5.rename(index=str_lower)) |
| |
| msg = "Either `index` or `columns` should be provided." |
| with self.assertRaisesRegex(ValueError, msg): |
| psdf1.rename() |
| msg = "`mapper` or `index` or `columns` should be either dict-like or function type." |
| with self.assertRaisesRegex(ValueError, msg): |
| psdf1.rename(mapper=[str_lower], axis=1) |
| msg = "Mapper dict should have the same value type." |
| with self.assertRaisesRegex(ValueError, msg): |
| psdf1.rename({"A": "a", "B": 2}, axis=1) |
| msg = r"level should be an integer between \[0, column_labels_level\)" |
| with self.assertRaisesRegex(ValueError, msg): |
| psdf2.rename(columns=str_lower, level=2) |
| msg = r"level should be an integer between \[0, 2\)" |
| with self.assertRaisesRegex(ValueError, msg): |
| psdf3.rename(index=str_lower, level=2) |
| |
| def test_rename_axis(self): |
| index = pd.Index(["A", "B", "C"], name="index") |
| columns = pd.Index(["numbers", "values"], name="cols") |
| pdf = pd.DataFrame([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], index=index, columns=columns) |
| psdf = ps.from_pandas(pdf) |
| |
| for axis in [0, "index"]: |
| self.assert_eq( |
| pdf.rename_axis("index2", axis=axis).sort_index(), |
| psdf.rename_axis("index2", axis=axis).sort_index(), |
| ) |
| self.assert_eq( |
| pdf.rename_axis(["index2"], axis=axis).sort_index(), |
| psdf.rename_axis(["index2"], axis=axis).sort_index(), |
| ) |
| |
| for axis in [1, "columns"]: |
| self.assert_eq( |
| pdf.rename_axis("cols2", axis=axis).sort_index(), |
| psdf.rename_axis("cols2", axis=axis).sort_index(), |
| ) |
| self.assert_eq( |
| pdf.rename_axis(["cols2"], axis=axis).sort_index(), |
| psdf.rename_axis(["cols2"], axis=axis).sort_index(), |
| ) |
| |
| pdf2 = pdf.copy() |
| psdf2 = psdf.copy() |
| pdf2.rename_axis("index2", axis="index", inplace=True) |
| psdf2.rename_axis("index2", axis="index", inplace=True) |
| self.assert_eq(pdf2.sort_index(), psdf2.sort_index()) |
| |
| self.assertRaises(ValueError, lambda: psdf.rename_axis(["index2", "index3"], axis=0)) |
| self.assertRaises(ValueError, lambda: psdf.rename_axis(["cols2", "cols3"], axis=1)) |
| self.assertRaises(TypeError, lambda: psdf.rename_axis(mapper=["index2"], index=["index3"])) |
| self.assertRaises(ValueError, lambda: psdf.rename_axis(ps)) |
| |
| self.assert_eq( |
| pdf.rename_axis(index={"index": "index2"}, columns={"cols": "cols2"}).sort_index(), |
| psdf.rename_axis(index={"index": "index2"}, columns={"cols": "cols2"}).sort_index(), |
| ) |
| |
| self.assert_eq( |
| pdf.rename_axis(index={"missing": "index2"}, columns={"missing": "cols2"}).sort_index(), |
| psdf.rename_axis( |
| index={"missing": "index2"}, columns={"missing": "cols2"} |
| ).sort_index(), |
| ) |
| |
| self.assert_eq( |
| pdf.rename_axis(index=str.upper, columns=str.upper).sort_index(), |
| psdf.rename_axis(index=str.upper, columns=str.upper).sort_index(), |
| ) |
| |
| index = pd.MultiIndex.from_tuples( |
| [("A", "B"), ("C", "D"), ("E", "F")], names=["index1", "index2"] |
| ) |
| columns = pd.MultiIndex.from_tuples( |
| [("numbers", "first"), ("values", "second")], names=["cols1", "cols2"] |
| ) |
| pdf = pd.DataFrame([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], index=index, columns=columns) |
| psdf = ps.from_pandas(pdf) |
| |
| for axis in [0, "index"]: |
| self.assert_eq( |
| pdf.rename_axis(["index3", "index4"], axis=axis).sort_index(), |
| psdf.rename_axis(["index3", "index4"], axis=axis).sort_index(), |
| ) |
| |
| for axis in [1, "columns"]: |
| self.assert_eq( |
| pdf.rename_axis(["cols3", "cols4"], axis=axis).sort_index(), |
| psdf.rename_axis(["cols3", "cols4"], axis=axis).sort_index(), |
| ) |
| |
| self.assertRaises( |
| ValueError, lambda: psdf.rename_axis(["index3", "index4", "index5"], axis=0) |
| ) |
| self.assertRaises(ValueError, lambda: psdf.rename_axis(["cols3", "cols4", "cols5"], axis=1)) |
| |
| self.assert_eq( |
| pdf.rename_axis(index={"index1": "index3"}, columns={"cols1": "cols3"}).sort_index(), |
| psdf.rename_axis(index={"index1": "index3"}, columns={"cols1": "cols3"}).sort_index(), |
| ) |
| |
| self.assert_eq( |
| pdf.rename_axis(index={"missing": "index3"}, columns={"missing": "cols3"}).sort_index(), |
| psdf.rename_axis( |
| index={"missing": "index3"}, columns={"missing": "cols3"} |
| ).sort_index(), |
| ) |
| |
| self.assert_eq( |
| pdf.rename_axis( |
| index={"index1": "index3", "index2": "index4"}, |
| columns={"cols1": "cols3", "cols2": "cols4"}, |
| ).sort_index(), |
| psdf.rename_axis( |
| index={"index1": "index3", "index2": "index4"}, |
| columns={"cols1": "cols3", "cols2": "cols4"}, |
| ).sort_index(), |
| ) |
| |
| self.assert_eq( |
| pdf.rename_axis(index=str.upper, columns=str.upper).sort_index(), |
| psdf.rename_axis(index=str.upper, columns=str.upper).sort_index(), |
| ) |
| |
| def test_multi_index_rename(self): |
| arrays = [[1, 1, 2, 2], ["red", "blue", "red", "blue"]] |
| idx = pd.MultiIndex.from_arrays(arrays, names=("number", "color")) |
| pdf = pd.DataFrame(np.random.randn(4, 5), idx) |
| psdf = ps.from_pandas(pdf) |
| |
| pmidx = pdf.index |
| psmidx = psdf.index |
| |
| self.assert_eq(psmidx.rename(["n", "c"]), pmidx.rename(["n", "c"])) |
| self.assert_eq(psdf.index.names, pdf.index.names) |
| |
| # non-string names |
| self.assert_eq(psmidx.rename([0, 1]), pmidx.rename([0, 1])) |
| self.assert_eq( |
| psmidx.rename([("x", "a"), ("y", "b")]), pmidx.rename([("x", "a"), ("y", "b")]) |
| ) |
| |
| psmidx.rename(["num", "col"], inplace=True) |
| pmidx.rename(["num", "col"], inplace=True) |
| |
| self.assert_eq(psmidx, pmidx) |
| self.assert_eq(psdf.index.names, pdf.index.names) |
| |
| self.assert_eq(psmidx.rename([None, None]), pmidx.rename([None, None])) |
| self.assert_eq(psdf.index.names, pdf.index.names) |
| |
| self.assertRaises(TypeError, lambda: psmidx.rename("number")) |
| self.assertRaises(TypeError, lambda: psmidx.rename(None)) |
| self.assertRaises(ValueError, lambda: psmidx.rename(["number"])) |
| |
| def test_multiindex_rename(self): |
| pidx = pd.MultiIndex.from_tuples([("a", "x", 1), ("b", "y", 2), ("c", "z", 3)]) |
| psidx = ps.from_pandas(pidx) |
| |
| pidx = pidx.rename(list("ABC")) |
| psidx = psidx.rename(list("ABC")) |
| self.assert_eq(pidx, psidx) |
| |
| pidx = pidx.rename(["my", "name", "is"]) |
| psidx = psidx.rename(["my", "name", "is"]) |
| self.assert_eq(pidx, psidx) |
| |
| def test_index_rename(self): |
| pdf = pd.DataFrame( |
| np.random.randn(10, 5), index=pd.Index([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], name="x") |
| ) |
| psdf = ps.from_pandas(pdf) |
| |
| pidx = pdf.index |
| psidx = psdf.index |
| |
| self.assert_eq(psidx.rename("y"), pidx.rename("y")) |
| self.assert_eq(psdf.index.names, pdf.index.names) |
| |
| # non-string names |
| self.assert_eq(psidx.rename(0), pidx.rename(0)) |
| self.assert_eq(psidx.rename(("y", 0)), pidx.rename(("y", 0))) |
| |
| psidx.rename("z", inplace=True) |
| pidx.rename("z", inplace=True) |
| |
| self.assert_eq(psidx, pidx) |
| self.assert_eq(psdf.index.names, pdf.index.names) |
| |
| self.assert_eq(psidx.rename(None), pidx.rename(None)) |
| self.assert_eq(psdf.index.names, pdf.index.names) |
| |
| self.assertRaises(TypeError, lambda: psidx.rename(["x", "y"])) |
| |
| |
| class FrameRenameTests( |
| FrameRenameMixin, |
| PandasOnSparkTestCase, |
| SQLTestUtils, |
| ): |
| pass |
| |
| |
| if __name__ == "__main__": |
| from pyspark.pandas.tests.indexes.test_rename 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) |
