# # 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 ComparisonTestBase from pyspark.testing.sqlutils import SQLTestUtils class FrameRenameMixin: @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=np.random.rand(9), ) @property def df_pair(self): pdf = self.pdf psdf = ps.from_pandas(pdf) return pdf, psdf 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(), ) class FrameRenameTests(FrameRenameMixin, ComparisonTestBase, 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)