| # |
| # 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 |
| |
| |
| # This file contains test cases for 'Missing data handling' |
| # https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/frame.html#missing-data-handling |
| class FrameMissingDataMixin: |
| def test_backfill(self): |
| pdf = pd.DataFrame( |
| { |
| "A": [None, 3, None, None], |
| "B": [2, 4, None, 3], |
| "C": [None, None, None, 1], |
| "D": [0, 1, 5, 4], |
| }, |
| columns=["A", "B", "C", "D"], |
| ) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(pdf.backfill(), psdf.backfill()) |
| |
| # Test `inplace=True` |
| pdf.backfill(inplace=True) |
| psdf.backfill(inplace=True) |
| self.assert_eq(pdf, psdf) |
| |
| def _test_dropna(self, pdf, axis): |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(psdf.dropna(axis=axis), pdf.dropna(axis=axis)) |
| self.assert_eq(psdf.dropna(axis=axis, how="all"), pdf.dropna(axis=axis, how="all")) |
| self.assert_eq(psdf.dropna(axis=axis, subset=["x"]), pdf.dropna(axis=axis, subset=["x"])) |
| self.assert_eq(psdf.dropna(axis=axis, subset="x"), pdf.dropna(axis=axis, subset=["x"])) |
| self.assert_eq( |
| psdf.dropna(axis=axis, subset=["y", "z"]), pdf.dropna(axis=axis, subset=["y", "z"]) |
| ) |
| self.assert_eq( |
| psdf.dropna(axis=axis, subset=["y", "z"], how="all"), |
| pdf.dropna(axis=axis, subset=["y", "z"], how="all"), |
| ) |
| |
| self.assert_eq(psdf.dropna(axis=axis, thresh=2), pdf.dropna(axis=axis, thresh=2)) |
| self.assert_eq( |
| psdf.dropna(axis=axis, thresh=1, subset=["y", "z"]), |
| pdf.dropna(axis=axis, thresh=1, subset=["y", "z"]), |
| ) |
| |
| pdf2 = pdf.copy() |
| psdf2 = psdf.copy() |
| pser = pdf2[pdf2.columns[0]] |
| psser = psdf2[psdf2.columns[0]] |
| pdf2.dropna(inplace=True, axis=axis) |
| psdf2.dropna(inplace=True, axis=axis) |
| self.assert_eq(psdf2, pdf2) |
| self.assert_eq(psser, pser) |
| |
| # multi-index |
| columns = pd.MultiIndex.from_tuples([("a", "x"), ("a", "y"), ("b", "z")]) |
| if axis == 0: |
| pdf.columns = columns |
| else: |
| pdf.index = columns |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(psdf.dropna(axis=axis), pdf.dropna(axis=axis)) |
| self.assert_eq(psdf.dropna(axis=axis, how="all"), pdf.dropna(axis=axis, how="all")) |
| self.assert_eq( |
| psdf.dropna(axis=axis, subset=[("a", "x")]), pdf.dropna(axis=axis, subset=[("a", "x")]) |
| ) |
| self.assert_eq( |
| psdf.dropna(axis=axis, subset=("a", "x")), pdf.dropna(axis=axis, subset=[("a", "x")]) |
| ) |
| self.assert_eq( |
| psdf.dropna(axis=axis, subset=[("a", "y"), ("b", "z")]), |
| pdf.dropna(axis=axis, subset=[("a", "y"), ("b", "z")]), |
| ) |
| self.assert_eq( |
| psdf.dropna(axis=axis, subset=[("a", "y"), ("b", "z")], how="all"), |
| pdf.dropna(axis=axis, subset=[("a", "y"), ("b", "z")], how="all"), |
| ) |
| |
| self.assert_eq(psdf.dropna(axis=axis, thresh=2), pdf.dropna(axis=axis, thresh=2)) |
| self.assert_eq( |
| psdf.dropna(axis=axis, thresh=1, subset=[("a", "y"), ("b", "z")]), |
| pdf.dropna(axis=axis, thresh=1, subset=[("a", "y"), ("b", "z")]), |
| ) |
| |
| def test_dropna_axis_index(self): |
| pdf = pd.DataFrame( |
| { |
| "x": [np.nan, 2, 3, 4, np.nan, 6], |
| "y": [1, 2, np.nan, 4, np.nan, np.nan], |
| "z": [1, 2, 3, 4, np.nan, np.nan], |
| }, |
| index=np.random.rand(6), |
| ) |
| |
| self._test_dropna(pdf, axis=0) |
| |
| # empty |
| pdf = pd.DataFrame(index=np.random.rand(6)) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(psdf.dropna(), pdf.dropna()) |
| self.assert_eq(psdf.dropna(how="all"), pdf.dropna(how="all")) |
| self.assert_eq(psdf.dropna(thresh=0), pdf.dropna(thresh=0)) |
| self.assert_eq(psdf.dropna(thresh=1), pdf.dropna(thresh=1)) |
| |
| # Only NA value |
| pdf["a"] = [np.nan] * 6 |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq(psdf.dropna(), pdf.dropna()) |
| self.assert_eq(psdf.dropna(how="all"), pdf.dropna(how="all")) |
| self.assert_eq(psdf.dropna(thresh=0), pdf.dropna(thresh=0)) |
| self.assert_eq(psdf.dropna(thresh=1), pdf.dropna(thresh=1)) |
| |
| with self.assertRaisesRegex(ValueError, "No axis named foo"): |
| psdf.dropna(axis="foo") |
| |
| self.assertRaises(KeyError, lambda: psdf.dropna(subset="1")) |
| with self.assertRaisesRegex(ValueError, "invalid how option: 1"): |
| psdf.dropna(how=1) |
| with self.assertRaisesRegex(TypeError, "must specify how or thresh"): |
| psdf.dropna(how=None) |
| |
| def test_dropna_axis_column(self): |
| pdf = pd.DataFrame( |
| { |
| "x": [np.nan, 2, 3, 4, np.nan, 6], |
| "y": [1, 2, np.nan, 4, np.nan, np.nan], |
| "z": [1, 2, 3, 4, np.nan, np.nan], |
| }, |
| index=[str(r) for r in np.random.rand(6)], |
| ).T |
| |
| self._test_dropna(pdf, axis=1) |
| |
| psdf = ps.from_pandas(pdf) |
| with self.assertRaisesRegex( |
| ValueError, "The length of each subset must be the same as the index size." |
| ): |
| psdf.dropna(subset=(["x", "y"]), axis=1) |
| |
| # empty |
| pdf = pd.DataFrame({"x": [], "y": [], "z": []}) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(psdf.dropna(axis=1), pdf.dropna(axis=1)) |
| self.assert_eq(psdf.dropna(axis=1, how="all"), pdf.dropna(axis=1, how="all")) |
| self.assert_eq(psdf.dropna(axis=1, thresh=0), pdf.dropna(axis=1, thresh=0)) |
| self.assert_eq(psdf.dropna(axis=1, thresh=1), pdf.dropna(axis=1, thresh=1)) |
| |
| def test_fillna(self): |
| pdf = pd.DataFrame( |
| { |
| "x": [np.nan, 2, 3, 4, np.nan, 6], |
| "y": [1, 2, np.nan, 4, np.nan, np.nan], |
| "z": [1, 2, 3, 4, np.nan, np.nan], |
| }, |
| index=np.random.rand(6), |
| ) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(psdf, pdf) |
| self.assert_eq(psdf.fillna(-1), pdf.fillna(-1)) |
| self.assert_eq( |
| psdf.fillna({"x": -1, "y": -2, "z": -5}), pdf.fillna({"x": -1, "y": -2, "z": -5}) |
| ) |
| self.assert_eq(pdf.fillna(method="ffill"), psdf.fillna(method="ffill")) |
| self.assert_eq(pdf.fillna(method="ffill", limit=2), psdf.fillna(method="ffill", limit=2)) |
| self.assert_eq(pdf.fillna(method="bfill"), psdf.fillna(method="bfill")) |
| self.assert_eq(pdf.fillna(method="bfill", limit=2), psdf.fillna(method="bfill", limit=2)) |
| |
| pdf = pdf.set_index(["x", "y"]) |
| psdf = ps.from_pandas(pdf) |
| # check multi index |
| self.assert_eq(psdf.fillna(-1), pdf.fillna(-1)) |
| self.assert_eq(pdf.fillna(method="bfill"), psdf.fillna(method="bfill")) |
| self.assert_eq(pdf.fillna(method="ffill"), psdf.fillna(method="ffill")) |
| |
| pser = pdf.z |
| psser = psdf.z |
| pdf.fillna({"x": -1, "y": -2, "z": -5}, inplace=True) |
| psdf.fillna({"x": -1, "y": -2, "z": -5}, inplace=True) |
| self.assert_eq(psdf, pdf) |
| self.assert_eq(psser, pser) |
| |
| pser = pdf.z |
| psser = psdf.z |
| pdf.fillna(0, inplace=True) |
| psdf.fillna(0, inplace=True) |
| self.assert_eq(psdf, pdf) |
| self.assert_eq(psser, pser) |
| |
| s_nan = pd.Series([-1, -2, -5], index=["x", "y", "z"], dtype=int) |
| self.assert_eq(psdf.fillna(s_nan), pdf.fillna(s_nan)) |
| |
| with self.assertRaisesRegex(NotImplementedError, "fillna currently only"): |
| psdf.fillna(-1, axis=1) |
| with self.assertRaisesRegex(NotImplementedError, "fillna currently only"): |
| psdf.fillna(-1, axis="columns") |
| with self.assertRaisesRegex(ValueError, "limit parameter for value is not support now"): |
| psdf.fillna(-1, limit=1) |
| with self.assertRaisesRegex(TypeError, "Unsupported.*DataFrame"): |
| psdf.fillna(pd.DataFrame({"x": [-1], "y": [-1], "z": [-1]})) |
| with self.assertRaisesRegex(TypeError, "Unsupported.*int64"): |
| psdf.fillna({"x": np.int64(-6), "y": np.int64(-4), "z": -5}) |
| with self.assertRaisesRegex(ValueError, "Expecting 'pad', 'ffill', 'backfill' or 'bfill'."): |
| psdf.fillna(method="xxx") |
| with self.assertRaisesRegex( |
| ValueError, "Must specify a fillna 'value' or 'method' parameter." |
| ): |
| psdf.fillna() |
| |
| # multi-index columns |
| pdf = pd.DataFrame( |
| { |
| ("x", "a"): [np.nan, 2, 3, 4, np.nan, 6], |
| ("x", "b"): [1, 2, np.nan, 4, np.nan, np.nan], |
| ("y", "c"): [1, 2, 3, 4, np.nan, np.nan], |
| }, |
| index=np.random.rand(6), |
| ) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(psdf.fillna(-1), pdf.fillna(-1)) |
| self.assert_eq( |
| psdf.fillna({("x", "a"): -1, ("x", "b"): -2, ("y", "c"): -5}), |
| pdf.fillna({("x", "a"): -1, ("x", "b"): -2, ("y", "c"): -5}), |
| ) |
| self.assert_eq(pdf.fillna(method="ffill"), psdf.fillna(method="ffill")) |
| self.assert_eq(pdf.fillna(method="ffill", limit=2), psdf.fillna(method="ffill", limit=2)) |
| self.assert_eq(pdf.fillna(method="bfill"), psdf.fillna(method="bfill")) |
| self.assert_eq(pdf.fillna(method="bfill", limit=2), psdf.fillna(method="bfill", limit=2)) |
| |
| self.assert_eq(psdf.fillna({"x": -1}), pdf.fillna({"x": -1})) |
| self.assert_eq( |
| psdf.fillna({"x": -1, ("x", "b"): -2}), pdf.fillna({"x": -1, ("x", "b"): -2}) |
| ) |
| self.assert_eq( |
| psdf.fillna({("x", "b"): -2, "x": -1}), pdf.fillna({("x", "b"): -2, "x": -1}) |
| ) |
| |
| # check multi index |
| pdf = pdf.set_index([("x", "a"), ("x", "b")]) |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq(psdf.fillna(-1), pdf.fillna(-1)) |
| self.assert_eq( |
| psdf.fillna({("x", "a"): -1, ("x", "b"): -2, ("y", "c"): -5}), |
| pdf.fillna({("x", "a"): -1, ("x", "b"): -2, ("y", "c"): -5}), |
| ) |
| |
| def test_replace(self): |
| pdf = pd.DataFrame( |
| { |
| "name": ["Ironman", "Captain America", "Thor", "Hulk"], |
| "weapon": ["Mark-45", "Shield", "Mjolnir", "Smash"], |
| }, |
| index=np.random.rand(4), |
| ) |
| psdf = ps.from_pandas(pdf) |
| |
| with self.assertRaisesRegex( |
| NotImplementedError, "replace currently works only for method='pad" |
| ): |
| psdf.replace(method="bfill") |
| with self.assertRaisesRegex( |
| NotImplementedError, "replace currently works only when limit=None" |
| ): |
| psdf.replace(limit=10) |
| with self.assertRaisesRegex( |
| NotImplementedError, "replace currently doesn't supports regex" |
| ): |
| psdf.replace(regex="") |
| |
| with self.assertRaisesRegex(ValueError, "Length of to_replace and value must be same"): |
| psdf.replace(to_replace=["Ironman"], value=["Spiderman", "Doctor Strange"]) |
| with self.assertRaisesRegex(TypeError, "Unsupported type function"): |
| psdf.replace("Ironman", lambda x: "Spiderman") |
| with self.assertRaisesRegex(TypeError, "Unsupported type function"): |
| psdf.replace(lambda x: "Ironman", "Spiderman") |
| |
| self.assert_eq(psdf.replace("Ironman", "Spiderman"), pdf.replace("Ironman", "Spiderman")) |
| self.assert_eq( |
| psdf.replace(["Ironman", "Captain America"], ["Rescue", "Hawkeye"]), |
| pdf.replace(["Ironman", "Captain America"], ["Rescue", "Hawkeye"]), |
| ) |
| self.assert_eq( |
| psdf.replace(("Ironman", "Captain America"), ("Rescue", "Hawkeye")), |
| pdf.replace(("Ironman", "Captain America"), ("Rescue", "Hawkeye")), |
| ) |
| |
| # inplace |
| pser = pdf.name |
| psser = psdf.name |
| pdf.replace("Ironman", "Spiderman", inplace=True) |
| psdf.replace("Ironman", "Spiderman", inplace=True) |
| self.assert_eq(psdf, pdf) |
| self.assert_eq(psser, pser) |
| |
| pdf = pd.DataFrame( |
| {"A": [0, 1, 2, 3, np.nan], "B": [5, 6, 7, 8, np.nan], "C": ["a", "b", "c", "d", None]}, |
| index=np.random.rand(5), |
| ) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(psdf.replace([0, 1, 2, 3, 5, 6], 4), pdf.replace([0, 1, 2, 3, 5, 6], 4)) |
| |
| self.assert_eq( |
| psdf.replace([0, 1, 2, 3, 5, 6], [6, 5, 4, 3, 2, 1]), |
| pdf.replace([0, 1, 2, 3, 5, 6], [6, 5, 4, 3, 2, 1]), |
| ) |
| |
| self.assert_eq(psdf.replace({0: 10, 1: 100, 7: 200}), pdf.replace({0: 10, 1: 100, 7: 200})) |
| |
| self.assert_eq( |
| psdf.replace({"A": [0, np.nan], "B": [5, np.nan]}, 100), |
| pdf.replace({"A": [0, np.nan], "B": [5, np.nan]}, 100), |
| ) |
| |
| self.assert_eq( |
| psdf.replace({"A": {0: 100, 4: 400, np.nan: 700}}), |
| pdf.replace({"A": {0: 100, 4: 400, np.nan: 700}}), |
| ) |
| self.assert_eq( |
| psdf.replace({"X": {0: 100, 4: 400, np.nan: 700}}), |
| pdf.replace({"X": {0: 100, 4: 400, np.nan: 700}}), |
| ) |
| |
| self.assert_eq(psdf.replace({"C": ["a", None]}, "e"), pdf.replace({"C": ["a", None]}, "e")) |
| |
| # multi-index columns |
| columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B"), ("Y", "C")]) |
| pdf.columns = columns |
| psdf.columns = columns |
| |
| self.assert_eq(psdf.replace([0, 1, 2, 3, 5, 6], 4), pdf.replace([0, 1, 2, 3, 5, 6], 4)) |
| |
| self.assert_eq( |
| psdf.replace([0, 1, 2, 3, 5, 6], [6, 5, 4, 3, 2, 1]), |
| pdf.replace([0, 1, 2, 3, 5, 6], [6, 5, 4, 3, 2, 1]), |
| ) |
| |
| self.assert_eq(psdf.replace({0: 10, 1: 100, 7: 200}), pdf.replace({0: 10, 1: 100, 7: 200})) |
| |
| self.assert_eq( |
| psdf.replace({("X", "A"): [0, np.nan], ("X", "B"): 5}, 100), |
| pdf.replace({("X", "A"): [0, np.nan], ("X", "B"): 5}, 100), |
| ) |
| |
| self.assert_eq( |
| psdf.replace({("X", "A"): {0: 100, 4: 400, np.nan: 700}}), |
| pdf.replace({("X", "A"): {0: 100, 4: 400, np.nan: 700}}), |
| ) |
| self.assert_eq( |
| psdf.replace({("X", "B"): {0: 100, 4: 400, np.nan: 700}}), |
| pdf.replace({("X", "B"): {0: 100, 4: 400, np.nan: 700}}), |
| ) |
| |
| self.assert_eq( |
| psdf.replace({("Y", "C"): ["a", None]}, "e"), |
| pdf.replace({("Y", "C"): ["a", None]}, "e"), |
| ) |
| |
| def test_ffill(self): |
| idx = np.random.rand(6) |
| pdf = pd.DataFrame( |
| { |
| "x": [np.nan, 2, 3, 4, np.nan, 6], |
| "y": [1, 2, np.nan, 4, np.nan, np.nan], |
| "z": [1, 2, 3, 4, np.nan, np.nan], |
| }, |
| index=idx, |
| ) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(psdf.ffill(), pdf.ffill()) |
| self.assert_eq(psdf.ffill(limit=1), pdf.ffill(limit=1)) |
| |
| pser = pdf.y |
| psser = psdf.y |
| |
| psdf.ffill(inplace=True) |
| pdf.ffill(inplace=True) |
| |
| self.assert_eq(psdf, pdf) |
| self.assert_eq(psser, pser) |
| self.assert_eq(psser[idx[2]], pser[idx[2]]) |
| |
| def test_bfill(self): |
| idx = np.random.rand(6) |
| pdf = pd.DataFrame( |
| { |
| "x": [np.nan, 2, 3, 4, np.nan, 6], |
| "y": [1, 2, np.nan, 4, np.nan, np.nan], |
| "z": [1, 2, 3, 4, np.nan, np.nan], |
| }, |
| index=idx, |
| ) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(psdf.bfill(), pdf.bfill()) |
| self.assert_eq(psdf.bfill(limit=1), pdf.bfill(limit=1)) |
| |
| pser = pdf.x |
| psser = psdf.x |
| |
| psdf.bfill(inplace=True) |
| pdf.bfill(inplace=True) |
| |
| self.assert_eq(psdf, pdf) |
| self.assert_eq(psser, pser) |
| self.assert_eq(psser[idx[0]], pser[idx[0]]) |
| |
| def test_pad(self): |
| pdf = pd.DataFrame( |
| { |
| "A": [None, 3, None, None], |
| "B": [2, 4, None, 3], |
| "C": [None, None, None, 1], |
| "D": [0, 1, 5, 4], |
| }, |
| columns=["A", "B", "C", "D"], |
| ) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(pdf.pad(), psdf.pad()) |
| |
| # Test `inplace=True` |
| pdf.pad(inplace=True) |
| psdf.pad(inplace=True) |
| self.assert_eq(pdf, psdf) |
| |
| |
| class FrameMissingDataTests( |
| FrameMissingDataMixin, |
| PandasOnSparkTestCase, |
| SQLTestUtils, |
| ): |
| pass |
| |
| |
| if __name__ == "__main__": |
| from pyspark.pandas.tests.computation.test_missing_data 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) |
