| # |
| # 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 SeriesIndexMixin: |
| def test_align(self): |
| pdf = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]}) |
| psdf = ps.from_pandas(pdf) |
| |
| for join in ["outer", "inner", "left", "right"]: |
| for axis in [None, 0]: |
| psser_l, psser_r = psdf.a.align(psdf.b, join=join, axis=axis) |
| pser_l, pser_r = pdf.a.align(pdf.b, join=join, axis=axis) |
| self.assert_eq(psser_l, pser_l) |
| self.assert_eq(psser_r, pser_r) |
| |
| psser_l, psdf_r = psdf.b.align(psdf[["b", "a"]], join=join, axis=axis) |
| pser_l, pdf_r = pdf.b.align(pdf[["b", "a"]], join=join, axis=axis) |
| self.assert_eq(psser_l, pser_l) |
| self.assert_eq(psdf_r, pdf_r) |
| |
| self.assertRaises(ValueError, lambda: psdf.a.align(psdf.b, axis=1)) |
| |
| def test_reindex(self): |
| index = ["A", "B", "C", "D", "E"] |
| pser = pd.Series([1.0, 2.0, 3.0, 4.0, None], index=index, name="x") |
| psser = ps.from_pandas(pser) |
| |
| self.assert_eq(pser, psser) |
| |
| self.assert_eq( |
| pser.reindex(["A", "B"]).sort_index(), |
| psser.reindex(["A", "B"]).sort_index(), |
| ) |
| |
| self.assert_eq( |
| pser.reindex(["A", "B", "2", "3"]).sort_index(), |
| psser.reindex(["A", "B", "2", "3"]).sort_index(), |
| ) |
| |
| self.assert_eq( |
| pser.reindex(["A", "E", "2"], fill_value=0).sort_index(), |
| psser.reindex(["A", "E", "2"], fill_value=0).sort_index(), |
| ) |
| |
| self.assertRaises(TypeError, lambda: psser.reindex(index=123)) |
| |
| def test_reindex_like(self): |
| data = [1.0, 2.0, None] |
| index = pd.Index(["A", "B", "C"], name="index1") |
| pser = pd.Series(data=data, index=index, name="name1") |
| psser = ps.from_pandas(pser) |
| |
| # Reindexing single Index on single Index |
| data2 = [3.0, None, 4.0] |
| index2 = pd.Index(["A", "C", "D"], name="index2") |
| pser2 = pd.Series(data=data2, index=index2, name="name2") |
| psser2 = ps.from_pandas(pser2) |
| |
| self.assert_eq( |
| pser.reindex_like(pser2).sort_index(), |
| psser.reindex_like(psser2).sort_index(), |
| ) |
| |
| self.assert_eq( |
| (pser + 1).reindex_like(pser2).sort_index(), |
| (psser + 1).reindex_like(psser2).sort_index(), |
| ) |
| |
| # Reindexing MultiIndex on single Index |
| index2 = pd.MultiIndex.from_tuples( |
| [("A", "G"), ("C", "D"), ("I", "J")], names=["index3", "index4"] |
| ) |
| pser2 = pd.Series(data=data2, index=index2, name="name2") |
| psser2 = ps.from_pandas(pser2) |
| |
| self.assert_eq( |
| pser.reindex_like(pser2).sort_index(), |
| psser.reindex_like(psser2).sort_index(), |
| ) |
| |
| self.assertRaises(TypeError, lambda: psser.reindex_like(index2)) |
| self.assertRaises(AssertionError, lambda: psser2.reindex_like(psser)) |
| |
| # Reindexing MultiIndex on MultiIndex |
| index = pd.MultiIndex.from_tuples( |
| [("A", "B"), ("C", "D"), ("E", "F")], names=["index1", "index2"] |
| ) |
| pser = pd.Series(data=data, index=index, name="name1") |
| psser = ps.from_pandas(pser) |
| |
| self.assert_eq( |
| pser.reindex_like(pser2).sort_index(), |
| psser.reindex_like(psser2).sort_index(), |
| ) |
| |
| # Reindexing with DataFrame |
| index2 = pd.MultiIndex.from_tuples( |
| [("A", "B"), ("C", "D"), ("E", "F")], names=["name3", "name4"] |
| ) |
| pdf = pd.DataFrame(data=data, index=index2) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq( |
| pser.reindex_like(pdf).sort_index(), |
| psser.reindex_like(psdf).sort_index(), |
| ) |
| |
| def test_reset_index(self): |
| pdf = pd.DataFrame({"foo": [1, 2, 3, 4]}, index=pd.Index(["a", "b", "c", "d"], name="idx")) |
| psdf = ps.from_pandas(pdf) |
| |
| pser = pdf.foo |
| psser = psdf.foo |
| |
| self.assert_eq(psser.reset_index(), pser.reset_index()) |
| self.assert_eq(psser.reset_index(name="values"), pser.reset_index(name="values")) |
| self.assert_eq(psser.reset_index(drop=True), pser.reset_index(drop=True)) |
| |
| # inplace |
| psser.reset_index(drop=True, inplace=True) |
| pser.reset_index(drop=True, inplace=True) |
| self.assert_eq(psser, pser) |
| self.assert_eq(psdf, pdf) |
| |
| def test_reset_index_with_default_index_types(self): |
| pser = pd.Series([1, 2, 3], name="0", index=np.random.rand(3)) |
| psser = ps.from_pandas(pser) |
| |
| with ps.option_context("compute.default_index_type", "sequence"): |
| self.assert_eq(psser.reset_index(), pser.reset_index()) |
| |
| with ps.option_context("compute.default_index_type", "distributed-sequence"): |
| # the order might be changed. |
| self.assert_eq(psser.reset_index().sort_index(), pser.reset_index()) |
| |
| with ps.option_context("compute.default_index_type", "distributed"): |
| # the index is different. |
| self.assert_eq( |
| psser.reset_index()._to_pandas().reset_index(drop=True), pser.reset_index() |
| ) |
| |
| def test_index_to_series_reset_index(self): |
| def check(psser, pser): |
| self.assert_eq(psser.reset_index(), pser.reset_index()) |
| self.assert_eq(psser.reset_index(drop=True), pser.reset_index(drop=True)) |
| |
| pser.reset_index(drop=True, inplace=True) |
| psser.reset_index(drop=True, inplace=True) |
| self.assert_eq(psser, pser) |
| |
| pdf = 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), |
| ) |
| psdf = ps.from_pandas(pdf) |
| check(psdf.index.to_series(), pdf.index.to_series()) |
| check(psdf.index.to_series(name="a"), pdf.index.to_series(name="a")) |
| check(psdf.index.to_series(name=("x", "a")), pdf.index.to_series(name=("x", "a"))) |
| |
| def test_idxmax(self): |
| pser = pd.Series(data=[1, 4, 5], index=["A", "B", "C"]) |
| psser = ps.Series(pser) |
| |
| self.assertEqual(psser.idxmax(), pser.idxmax()) |
| self.assertEqual(psser.idxmax(skipna=False), pser.idxmax(skipna=False)) |
| |
| index = pd.MultiIndex.from_arrays( |
| [["a", "a", "b", "b"], ["c", "d", "e", "f"]], names=("first", "second") |
| ) |
| pser = pd.Series(data=[1, 2, 4, 5], index=index) |
| psser = ps.Series(pser) |
| |
| self.assertEqual(psser.idxmax(), pser.idxmax()) |
| self.assertEqual(psser.idxmax(skipna=False), pser.idxmax(skipna=False)) |
| |
| psser = ps.Series([]) |
| with self.assertRaisesRegex(ValueError, "an empty sequence"): |
| psser.idxmax() |
| |
| pser = pd.Series([1, 100, None, 100, 1, 100], index=[10, 3, 5, 2, 1, 8]) |
| psser = ps.Series(pser) |
| |
| self.assertEqual(psser.idxmax(), pser.idxmax()) |
| self.assertEqual(repr(psser.idxmax(skipna=False)), repr(pser.idxmax(skipna=False))) |
| |
| def test_idxmin(self): |
| pser = pd.Series(data=[1, 4, 5], index=["A", "B", "C"]) |
| psser = ps.Series(pser) |
| |
| self.assertEqual(psser.idxmin(), pser.idxmin()) |
| self.assertEqual(psser.idxmin(skipna=False), pser.idxmin(skipna=False)) |
| |
| index = pd.MultiIndex.from_arrays( |
| [["a", "a", "b", "b"], ["c", "d", "e", "f"]], names=("first", "second") |
| ) |
| pser = pd.Series(data=[1, 2, 4, 5], index=index) |
| psser = ps.Series(pser) |
| |
| self.assertEqual(psser.idxmin(), pser.idxmin()) |
| self.assertEqual(psser.idxmin(skipna=False), pser.idxmin(skipna=False)) |
| |
| psser = ps.Series([]) |
| with self.assertRaisesRegex(ValueError, "an empty sequence"): |
| psser.idxmin() |
| |
| pser = pd.Series([1, 100, None, 100, 1, 100], index=[10, 3, 5, 2, 1, 8]) |
| psser = ps.Series(pser) |
| |
| self.assertEqual(psser.idxmin(), pser.idxmin()) |
| self.assertEqual(repr(psser.idxmin(skipna=False)), repr(pser.idxmin(skipna=False))) |
| |
| def test_index(self): |
| # to check setting name of Index properly. |
| idx = pd.Index([1, 2, 3, 4, 5, 6, 7, 8, 9]) |
| pser = pd.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], index=idx) |
| psser = ps.from_pandas(pser) |
| |
| psser.name = "koalas" |
| pser.name = "koalas" |
| self.assert_eq(psser.index.name, pser.index.name) |
| |
| # for check setting names of MultiIndex properly. |
| psser.names = ["hello", "koalas"] |
| pser.names = ["hello", "koalas"] |
| self.assert_eq(psser.index.names, pser.index.names) |
| |
| def test_axes(self): |
| pser = pd.Series([90, 91, 85], index=[2, 4, 1]) |
| psser = ps.from_pandas(pser) |
| self.assert_eq(psser.axes, pser.axes) |
| |
| # for MultiIndex |
| midx = pd.MultiIndex( |
| [["lama", "cow", "falcon"], ["speed", "weight", "length"]], |
| [[0, 0, 0, 1, 1, 1, 2, 2, 2], [0, 1, 2, 0, 1, 2, 0, 1, 2]], |
| ) |
| pser = pd.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], index=midx) |
| psser = ps.from_pandas(pser) |
| self.assert_eq(psser.axes, pser.axes) |
| |
| def test_swaplevel(self): |
| # MultiIndex with two levels |
| arrays = [[1, 1, 2, 2], ["red", "blue", "red", "blue"]] |
| pidx = pd.MultiIndex.from_arrays(arrays, names=("number", "color")) |
| pser = pd.Series(["a", "b", "c", "d"], index=pidx) |
| psser = ps.from_pandas(pser) |
| self.assert_eq(pser.swaplevel(), psser.swaplevel()) |
| self.assert_eq(pser.swaplevel(0, 1), psser.swaplevel(0, 1)) |
| self.assert_eq(pser.swaplevel(1, 1), psser.swaplevel(1, 1)) |
| self.assert_eq(pser.swaplevel("number", "color"), psser.swaplevel("number", "color")) |
| |
| # MultiIndex with more than two levels |
| arrays = [[1, 1, 2, 2], ["red", "blue", "red", "blue"], ["l", "m", "s", "xs"]] |
| pidx = pd.MultiIndex.from_arrays(arrays, names=("number", "color", "size")) |
| pser = pd.Series(["a", "b", "c", "d"], index=pidx) |
| psser = ps.from_pandas(pser) |
| self.assert_eq(pser.swaplevel(), psser.swaplevel()) |
| self.assert_eq(pser.swaplevel(0, 1), psser.swaplevel(0, 1)) |
| self.assert_eq(pser.swaplevel(0, 2), psser.swaplevel(0, 2)) |
| self.assert_eq(pser.swaplevel(1, 2), psser.swaplevel(1, 2)) |
| self.assert_eq(pser.swaplevel(1, 1), psser.swaplevel(1, 1)) |
| self.assert_eq(pser.swaplevel(-1, -2), psser.swaplevel(-1, -2)) |
| self.assert_eq(pser.swaplevel("number", "color"), psser.swaplevel("number", "color")) |
| self.assert_eq(pser.swaplevel("number", "size"), psser.swaplevel("number", "size")) |
| self.assert_eq(pser.swaplevel("color", "size"), psser.swaplevel("color", "size")) |
| |
| # Error conditions |
| self.assertRaises(AssertionError, lambda: ps.Series([1, 2]).swaplevel()) |
| self.assertRaises(IndexError, lambda: psser.swaplevel(0, 9)) |
| self.assertRaises(KeyError, lambda: psser.swaplevel("not_number", "color")) |
| self.assertRaises(AssertionError, lambda: psser.swaplevel(copy=False)) |
| |
| def test_swapaxes(self): |
| pser = pd.Series([1, 2, 3], index=["x", "y", "z"], name="ser") |
| psser = ps.from_pandas(pser) |
| |
| self.assert_eq(psser.swapaxes(0, 0), pser.swapaxes(0, 0)) |
| self.assert_eq(psser.swapaxes("index", "index"), pser.swapaxes("index", "index")) |
| self.assert_eq((psser + 1).swapaxes(0, 0), (pser + 1).swapaxes(0, 0)) |
| |
| self.assertRaises(AssertionError, lambda: psser.swapaxes(0, 1, copy=False)) |
| self.assertRaises(ValueError, lambda: psser.swapaxes(0, 1)) |
| self.assertRaises(ValueError, lambda: psser.swapaxes("index", "columns")) |
| |
| def test_droplevel(self): |
| pser = pd.Series( |
| [1, 2, 3], |
| index=pd.MultiIndex.from_tuples( |
| [("x", "a", "q"), ("x", "b", "w"), ("y", "c", "e")], |
| names=["level_1", "level_2", "level_3"], |
| ), |
| ) |
| psser = ps.from_pandas(pser) |
| |
| self.assert_eq(pser.droplevel(0), psser.droplevel(0)) |
| self.assert_eq(pser.droplevel("level_1"), psser.droplevel("level_1")) |
| self.assert_eq(pser.droplevel(-1), psser.droplevel(-1)) |
| self.assert_eq(pser.droplevel([0]), psser.droplevel([0])) |
| self.assert_eq(pser.droplevel(["level_1"]), psser.droplevel(["level_1"])) |
| self.assert_eq(pser.droplevel((0,)), psser.droplevel((0,))) |
| self.assert_eq(pser.droplevel(("level_1",)), psser.droplevel(("level_1",))) |
| self.assert_eq(pser.droplevel([0, 2]), psser.droplevel([0, 2])) |
| self.assert_eq( |
| pser.droplevel(["level_1", "level_3"]), psser.droplevel(["level_1", "level_3"]) |
| ) |
| self.assert_eq(pser.droplevel((1, 2)), psser.droplevel((1, 2))) |
| self.assert_eq( |
| pser.droplevel(("level_2", "level_3")), psser.droplevel(("level_2", "level_3")) |
| ) |
| |
| with self.assertRaisesRegex(KeyError, "Level {0, 1, 2} not found"): |
| psser.droplevel({0, 1, 2}) |
| with self.assertRaisesRegex(KeyError, "Level level_100 not found"): |
| psser.droplevel(["level_1", "level_100"]) |
| with self.assertRaisesRegex(IndexError, "Too many levels: Index has only 3 levels, not 11"): |
| psser.droplevel(10) |
| with self.assertRaisesRegex( |
| IndexError, |
| "Too many levels: Index has only 3 levels, -10 is not a valid level number", |
| ): |
| psser.droplevel(-10) |
| with self.assertRaisesRegex( |
| ValueError, |
| "Cannot remove 3 levels from an index with 3 levels: " |
| "at least one level must be left.", |
| ): |
| psser.droplevel([0, 1, 2]) |
| with self.assertRaisesRegex( |
| ValueError, |
| "Cannot remove 5 levels from an index with 3 levels: " |
| "at least one level must be left.", |
| ): |
| psser.droplevel([1, 1, 1, 1, 1]) |
| |
| # Tupled names |
| pser.index.names = [("a", "1"), ("b", "2"), ("c", "3")] |
| psser = ps.from_pandas(pser) |
| |
| self.assert_eq( |
| pser.droplevel([("a", "1"), ("c", "3")]), psser.droplevel([("a", "1"), ("c", "3")]) |
| ) |
| |
| def test_last_valid_index(self): |
| pser = pd.Series([250, 1.5, 320, 1, 0.3, None, None, None, None]) |
| psser = ps.from_pandas(pser) |
| self.assert_eq(pser.last_valid_index(), psser.last_valid_index()) |
| |
| # MultiIndex columns |
| midx = pd.MultiIndex( |
| [["lama", "cow", "falcon"], ["speed", "weight", "length"]], |
| [[0, 0, 0, 1, 1, 1, 2, 2, 2], [0, 1, 2, 0, 1, 2, 0, 1, 2]], |
| ) |
| pser.index = midx |
| psser = ps.from_pandas(pser) |
| self.assert_eq(pser.last_valid_index(), psser.last_valid_index()) |
| |
| # Empty Series |
| pser = pd.Series([]) |
| psser = ps.from_pandas(pser) |
| self.assert_eq(pser.last_valid_index(), psser.last_valid_index()) |
| |
| def test_first_valid_index(self): |
| # Empty Series |
| pser = pd.Series([]) |
| psser = ps.from_pandas(pser) |
| self.assert_eq(pser.first_valid_index(), psser.first_valid_index()) |
| |
| def test_rename_axis(self): |
| index = pd.Index(["A", "B", "C"], name="index") |
| pdf = pd.DataFrame({"x": [1.0, 2.0, 3.0]}, index=index) |
| psdf = ps.from_pandas(pdf) |
| pser, psser = pdf.x, psdf.x |
| |
| self.assert_eq( |
| pser.rename_axis("index2").sort_index(), |
| psser.rename_axis("index2").sort_index(), |
| ) |
| |
| self.assert_eq( |
| (pser + 1).rename_axis("index2").sort_index(), |
| (psser + 1).rename_axis("index2").sort_index(), |
| ) |
| |
| self.assertRaises(ValueError, lambda: psser.rename_axis(["index2", "index3"])) |
| self.assertRaises(TypeError, lambda: psser.rename_axis(mapper=["index2"], index=["index3"])) |
| |
| self.assert_eq( |
| pser.rename_axis(index={"index": "index2", "missing": "index4"}).sort_index(), |
| psser.rename_axis(index={"index": "index2", "missing": "index4"}).sort_index(), |
| ) |
| |
| self.assert_eq( |
| pser.rename_axis(index=str.upper).sort_index(), |
| psser.rename_axis(index=str.upper).sort_index(), |
| ) |
| |
| pser.rename_axis("index2", inplace=True) |
| psser.rename_axis("index2", inplace=True) |
| self.assert_eq(pser.sort_index(), psser.sort_index()) |
| # Note: in pandas, pdf.x's index is renamed, whereas pdf's index isn't due to a bug. |
| self.assert_eq(pdf, psdf) |
| |
| index = pd.MultiIndex.from_tuples( |
| [("A", "B"), ("C", "D"), ("E", "F")], names=["index1", "index2"] |
| ) |
| pser = pd.Series([1.0, 2.0, 3.0], index=index, name="name") |
| psser = ps.from_pandas(pser) |
| |
| self.assert_eq( |
| pser.rename_axis(["index3", "index4"]).sort_index(), |
| psser.rename_axis(["index3", "index4"]).sort_index(), |
| ) |
| |
| self.assertRaises(ValueError, lambda: psser.rename_axis(["index3", "index4", "index5"])) |
| |
| self.assert_eq( |
| pser.rename_axis( |
| index={"index1": "index3", "index2": "index4", "missing": "index5"} |
| ).sort_index(), |
| psser.rename_axis( |
| index={"index1": "index3", "index2": "index4", "missing": "index5"} |
| ).sort_index(), |
| ) |
| |
| self.assert_eq( |
| pser.rename_axis(index=str.upper).sort_index(), |
| psser.rename_axis(index=str.upper).sort_index(), |
| ) |
| |
| |
| class SeriesIndexTests( |
| SeriesIndexMixin, |
| PandasOnSparkTestCase, |
| SQLTestUtils, |
| ): |
| pass |
| |
| |
| if __name__ == "__main__": |
| from pyspark.pandas.tests.series.test_index 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) |
