| # |
| # 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 pandas.tseries.offsets import DateOffset |
| |
| from pyspark import pandas as ps |
| from pyspark.errors import PySparkValueError |
| from pyspark.pandas.config import option_context |
| from pyspark.testing.pandasutils import PandasOnSparkTestCase |
| from pyspark.testing.sqlutils import SQLTestUtils |
| |
| |
| # This file contains test cases for 'Reindexing / Selection / Label manipulation' |
| # https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/frame.html#reindexing-selection-label-manipulation |
| class FrameReindexingMixin: |
| def test_add_prefix(self): |
| pdf = pd.DataFrame({"A": [1, 2, 3, 4], "B": [3, 4, 5, 6]}, index=np.random.rand(4)) |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq(pdf.add_prefix("col_"), psdf.add_prefix("col_")) |
| |
| columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B")]) |
| pdf.columns = columns |
| psdf.columns = columns |
| self.assert_eq(pdf.add_prefix("col_"), psdf.add_prefix("col_")) |
| |
| def test_add_suffix(self): |
| pdf = pd.DataFrame({"A": [1, 2, 3, 4], "B": [3, 4, 5, 6]}, index=np.random.rand(4)) |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq(pdf.add_suffix("first_series"), psdf.add_suffix("first_series")) |
| |
| columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B")]) |
| pdf.columns = columns |
| psdf.columns = columns |
| self.assert_eq(pdf.add_suffix("first_series"), psdf.add_suffix("first_series")) |
| |
| def test_at_time(self): |
| idx = pd.date_range("2018-04-09", periods=4, freq="1D20min") |
| pdf = pd.DataFrame({"A": [1, 2, 3, 4]}, index=idx) |
| psdf = ps.from_pandas(pdf) |
| psdf.at_time("0:20") |
| self.assert_eq( |
| pdf.at_time("0:20").sort_index(), |
| psdf.at_time("0:20").sort_index(), |
| ) |
| |
| # Index name is 'ts' |
| pdf.index.name = "ts" |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq( |
| pdf.at_time("0:20").sort_index(), |
| psdf.at_time("0:20").sort_index(), |
| ) |
| |
| # Index name is 'ts', column label is 'index' |
| pdf.columns = pd.Index(["index"]) |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq( |
| pdf.at_time("0:40").sort_index(), |
| psdf.at_time("0:40").sort_index(), |
| ) |
| |
| # Both index name and column label are 'index' |
| pdf.index.name = "index" |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq( |
| pdf.at_time("0:40").sort_index(), |
| psdf.at_time("0:40").sort_index(), |
| ) |
| |
| # Index name is 'index', column label is ('X', 'A') |
| pdf.columns = pd.MultiIndex.from_arrays([["X"], ["A"]]) |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq( |
| pdf.at_time("0:40").sort_index(), |
| psdf.at_time("0:40").sort_index(), |
| ) |
| |
| with self.assertRaisesRegex(NotImplementedError, "'asof' argument is not supported"): |
| psdf.at_time("0:15", asof=True) |
| |
| with self.assertRaisesRegex(NotImplementedError, "at_time currently only works for axis=0"): |
| psdf.at_time("0:15", axis=1) |
| |
| psdf = ps.DataFrame({"A": [1, 2, 3, 4]}) |
| with self.assertRaisesRegex(TypeError, "Index must be DatetimeIndex"): |
| psdf.at_time("0:15") |
| |
| def test_between_time(self): |
| idx = pd.date_range("2018-04-09", periods=4, freq="1D20min") |
| pdf = pd.DataFrame({"A": [1, 2, 3, 4]}, index=idx) |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq( |
| pdf.between_time("0:15", "0:45").sort_index(), |
| psdf.between_time("0:15", "0:45").sort_index(), |
| ) |
| |
| pdf.index.name = "ts" |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq( |
| pdf.between_time("0:15", "0:45").sort_index(), |
| psdf.between_time("0:15", "0:45").sort_index(), |
| ) |
| |
| # Column label is 'index' |
| pdf.columns = pd.Index(["index"]) |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq( |
| pdf.between_time("0:15", "0:45").sort_index(), |
| psdf.between_time("0:15", "0:45").sort_index(), |
| ) |
| |
| # Both index name and column label are 'index' |
| pdf.index.name = "index" |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq( |
| pdf.between_time("0:15", "0:45").sort_index(), |
| psdf.between_time("0:15", "0:45").sort_index(), |
| ) |
| |
| # Index name is 'index', column label is ('X', 'A') |
| pdf.columns = pd.MultiIndex.from_arrays([["X"], ["A"]]) |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq( |
| pdf.between_time("0:15", "0:45").sort_index(), |
| psdf.between_time("0:15", "0:45").sort_index(), |
| ) |
| |
| with self.assertRaisesRegex( |
| NotImplementedError, "between_time currently only works for axis=0" |
| ): |
| psdf.between_time("0:15", "0:45", axis=1) |
| |
| psdf = ps.DataFrame({"A": [1, 2, 3, 4]}) |
| with self.assertRaisesRegex(TypeError, "Index must be DatetimeIndex"): |
| psdf.between_time("0:15", "0:45") |
| |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq( |
| pdf.between_time("0:15", "0:45", inclusive="neither").sort_index(), |
| psdf.between_time("0:15", "0:45", inclusive="neither").sort_index(), |
| ) |
| |
| self.assert_eq( |
| pdf.between_time("0:15", "0:45", inclusive="left").sort_index(), |
| psdf.between_time("0:15", "0:45", inclusive="left").sort_index(), |
| ) |
| |
| self.assert_eq( |
| pdf.between_time("0:15", "0:45", inclusive="right").sort_index(), |
| psdf.between_time("0:15", "0:45", inclusive="right").sort_index(), |
| ) |
| |
| with self.assertRaises(PySparkValueError) as ctx: |
| psdf.between_time("0:15", "0:45", inclusive="") |
| |
| self.check_error( |
| exception=ctx.exception, |
| errorClass="VALUE_NOT_ALLOWED", |
| messageParameters={ |
| "arg_name": "inclusive", |
| "allowed_values": str(["left", "right", "both", "neither"]), |
| }, |
| ) |
| |
| def test_drop(self): |
| pdf = pd.DataFrame({"x": [1, 2], "y": [3, 4], "z": [5, 6]}, index=np.random.rand(2)) |
| psdf = ps.from_pandas(pdf) |
| |
| # Assert 'labels' or 'columns' parameter is set |
| expected_error_message = "Need to specify at least one of 'labels' or 'columns'" |
| with self.assertRaisesRegex(ValueError, expected_error_message): |
| psdf.drop() |
| |
| # |
| # Drop columns |
| # |
| |
| # Assert using a str for 'labels' works |
| self.assert_eq(psdf.drop("x", axis=1), pdf.drop("x", axis=1)) |
| self.assert_eq((psdf + 1).drop("x", axis=1), (pdf + 1).drop("x", axis=1)) |
| # Assert using a list for 'labels' works |
| self.assert_eq(psdf.drop(["y", "z"], axis=1), pdf.drop(["y", "z"], axis=1)) |
| self.assert_eq(psdf.drop(["x", "y", "z"], axis=1), pdf.drop(["x", "y", "z"], axis=1)) |
| # Assert using 'columns' instead of 'labels' produces the same results |
| self.assert_eq(psdf.drop(columns="x"), pdf.drop(columns="x")) |
| self.assert_eq(psdf.drop(columns=["y", "z"]), pdf.drop(columns=["y", "z"])) |
| self.assert_eq(psdf.drop(columns=["x", "y", "z"]), pdf.drop(columns=["x", "y", "z"])) |
| self.assert_eq(psdf.drop(columns=[]), pdf.drop(columns=[])) |
| |
| columns = pd.MultiIndex.from_tuples([(1, "x"), (1, "y"), (2, "z")]) |
| pdf.columns = columns |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(psdf.drop(columns=1), pdf.drop(columns=1)) |
| self.assert_eq(psdf.drop(columns=(1, "x")), pdf.drop(columns=(1, "x"))) |
| self.assert_eq(psdf.drop(columns=[(1, "x"), 2]), pdf.drop(columns=[(1, "x"), 2])) |
| self.assert_eq( |
| psdf.drop(columns=[(1, "x"), (1, "y"), (2, "z")]), |
| pdf.drop(columns=[(1, "x"), (1, "y"), (2, "z")]), |
| ) |
| |
| self.assertRaises(KeyError, lambda: psdf.drop(columns=3)) |
| self.assertRaises(KeyError, lambda: psdf.drop(columns=(1, "z"))) |
| |
| pdf.index = pd.MultiIndex.from_tuples([("i", 0), ("j", 1)]) |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq( |
| psdf.drop(columns=[(1, "x"), (1, "y"), (2, "z")]), |
| pdf.drop(columns=[(1, "x"), (1, "y"), (2, "z")]), |
| ) |
| |
| # non-string names |
| pdf = pd.DataFrame({10: [1, 2], 20: [3, 4], 30: [5, 6]}, index=np.random.rand(2)) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(psdf.drop(10, axis=1), pdf.drop(10, axis=1)) |
| self.assert_eq(psdf.drop([20, 30], axis=1), pdf.drop([20, 30], axis=1)) |
| |
| # |
| # Drop rows |
| # |
| |
| pdf = pd.DataFrame({"X": [1, 2, 3], "Y": [4, 5, 6], "Z": [7, 8, 9]}, index=["A", "B", "C"]) |
| psdf = ps.from_pandas(pdf) |
| |
| # Given labels (and axis = 0) |
| self.assert_eq(psdf.drop(labels="A", axis=0), pdf.drop(labels="A", axis=0)) |
| self.assert_eq(psdf.drop(labels="A"), pdf.drop(labels="A")) |
| self.assert_eq((psdf + 1).drop(labels="A"), (pdf + 1).drop(labels="A")) |
| self.assert_eq(psdf.drop(labels=["A", "C"], axis=0), pdf.drop(labels=["A", "C"], axis=0)) |
| self.assert_eq( |
| psdf.drop(labels=["A", "B", "C"], axis=0), pdf.drop(labels=["A", "B", "C"], axis=0) |
| ) |
| |
| with ps.option_context("compute.isin_limit", 2): |
| self.assert_eq( |
| psdf.drop(labels=["A", "B", "C"], axis=0), pdf.drop(labels=["A", "B", "C"], axis=0) |
| ) |
| |
| # Given index |
| self.assert_eq(psdf.drop(index="A"), pdf.drop(index="A")) |
| self.assert_eq(psdf.drop(index=["A", "C"]), pdf.drop(index=["A", "C"])) |
| self.assert_eq(psdf.drop(index=["A", "B", "C"]), pdf.drop(index=["A", "B", "C"])) |
| self.assert_eq(psdf.drop(index=[]), pdf.drop(index=[])) |
| |
| with ps.option_context("compute.isin_limit", 2): |
| self.assert_eq(psdf.drop(index=["A", "B", "C"]), pdf.drop(index=["A", "B", "C"])) |
| |
| # Non-string names |
| pdf.index = [10, 20, 30] |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq(psdf.drop(labels=10, axis=0), pdf.drop(labels=10, axis=0)) |
| self.assert_eq(psdf.drop(labels=[10, 30], axis=0), pdf.drop(labels=[10, 30], axis=0)) |
| self.assert_eq( |
| psdf.drop(labels=[10, 20, 30], axis=0), pdf.drop(labels=[10, 20, 30], axis=0) |
| ) |
| |
| with ps.option_context("compute.isin_limit", 2): |
| self.assert_eq( |
| psdf.drop(labels=[10, 20, 30], axis=0), pdf.drop(labels=[10, 20, 30], axis=0) |
| ) |
| |
| # MultiIndex |
| pdf.index = pd.MultiIndex.from_tuples([("a", "x"), ("b", "y"), ("c", "z")]) |
| psdf = ps.from_pandas(pdf) |
| self.assertRaises(NotImplementedError, lambda: psdf.drop(labels=[("a", "x")])) |
| |
| # |
| # Drop rows and columns |
| # |
| pdf = pd.DataFrame({"X": [1, 2, 3], "Y": [4, 5, 6], "Z": [7, 8, 9]}, index=["A", "B", "C"]) |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq(psdf.drop(index="A", columns="X"), pdf.drop(index="A", columns="X")) |
| self.assert_eq( |
| psdf.drop(index=["A", "C"], columns=["X", "Z"]), |
| pdf.drop(index=["A", "C"], columns=["X", "Z"]), |
| ) |
| self.assert_eq( |
| psdf.drop(index=["A", "B", "C"], columns=["X", "Z"]), |
| pdf.drop(index=["A", "B", "C"], columns=["X", "Z"]), |
| ) |
| with ps.option_context("compute.isin_limit", 2): |
| self.assert_eq( |
| psdf.drop(index=["A", "B", "C"], columns=["X", "Z"]), |
| pdf.drop(index=["A", "B", "C"], columns=["X", "Z"]), |
| ) |
| self.assert_eq( |
| psdf.drop(index=[], columns=["X", "Z"]), |
| pdf.drop(index=[], columns=["X", "Z"]), |
| ) |
| self.assert_eq( |
| psdf.drop(index=["A", "B", "C"], columns=[]), |
| pdf.drop(index=["A", "B", "C"], columns=[]), |
| ) |
| self.assert_eq( |
| psdf.drop(index=[], columns=[]), |
| pdf.drop(index=[], columns=[]), |
| ) |
| self.assertRaises( |
| ValueError, |
| lambda: psdf.drop(labels="A", axis=0, columns="X"), |
| ) |
| |
| def test_droplevel(self): |
| pdf = ( |
| pd.DataFrame([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]) |
| .set_index([0, 1]) |
| .rename_axis(["a", "b"]) |
| ) |
| pdf.columns = pd.MultiIndex.from_tuples( |
| [("c", "e"), ("d", "f")], names=["level_1", "level_2"] |
| ) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assertRaises(ValueError, lambda: psdf.droplevel(["a", "b"])) |
| self.assertRaises(ValueError, lambda: psdf.droplevel([1, 1, 1, 1, 1])) |
| self.assertRaises(IndexError, lambda: psdf.droplevel(2)) |
| self.assertRaises(IndexError, lambda: psdf.droplevel(-3)) |
| self.assertRaises(KeyError, lambda: psdf.droplevel({"a"})) |
| self.assertRaises(KeyError, lambda: psdf.droplevel({"a": 1})) |
| |
| self.assertRaises(ValueError, lambda: psdf.droplevel(["level_1", "level_2"], axis=1)) |
| self.assertRaises(IndexError, lambda: psdf.droplevel(2, axis=1)) |
| self.assertRaises(IndexError, lambda: psdf.droplevel(-3, axis=1)) |
| self.assertRaises(KeyError, lambda: psdf.droplevel({"level_1"}, axis=1)) |
| self.assertRaises(KeyError, lambda: psdf.droplevel({"level_1": 1}, axis=1)) |
| |
| self.assert_eq(pdf.droplevel("a"), psdf.droplevel("a")) |
| self.assert_eq(pdf.droplevel(["a"]), psdf.droplevel(["a"])) |
| self.assert_eq(pdf.droplevel(("a",)), psdf.droplevel(("a",))) |
| self.assert_eq(pdf.droplevel(0), psdf.droplevel(0)) |
| self.assert_eq(pdf.droplevel(-1), psdf.droplevel(-1)) |
| |
| self.assert_eq(pdf.droplevel("level_1", axis=1), psdf.droplevel("level_1", axis=1)) |
| self.assert_eq(pdf.droplevel(["level_1"], axis=1), psdf.droplevel(["level_1"], axis=1)) |
| self.assert_eq(pdf.droplevel(("level_1",), axis=1), psdf.droplevel(("level_1",), axis=1)) |
| self.assert_eq(pdf.droplevel(0, axis=1), psdf.droplevel(0, axis=1)) |
| self.assert_eq(pdf.droplevel(-1, axis=1), psdf.droplevel(-1, axis=1)) |
| |
| # Tupled names |
| pdf.columns.names = [("level", 1), ("level", 2)] |
| pdf.index.names = [("a", 10), ("x", 20)] |
| psdf = ps.from_pandas(pdf) |
| |
| self.assertRaises(KeyError, lambda: psdf.droplevel("a")) |
| self.assertRaises(KeyError, lambda: psdf.droplevel(("a", 10))) |
| |
| self.assert_eq(pdf.droplevel([("a", 10)]), psdf.droplevel([("a", 10)])) |
| self.assert_eq( |
| pdf.droplevel([("level", 1)], axis=1), psdf.droplevel([("level", 1)], axis=1) |
| ) |
| |
| # non-string names |
| pdf = ( |
| pd.DataFrame([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]) |
| .set_index([0, 1]) |
| .rename_axis([10.0, 20.0]) |
| ) |
| pdf.columns = pd.MultiIndex.from_tuples([("c", "e"), ("d", "f")], names=[100.0, 200.0]) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(pdf.droplevel(10.0), psdf.droplevel(10.0)) |
| self.assert_eq(pdf.droplevel([10.0]), psdf.droplevel([10.0])) |
| self.assert_eq(pdf.droplevel((10.0,)), psdf.droplevel((10.0,))) |
| self.assert_eq(pdf.droplevel(0), psdf.droplevel(0)) |
| self.assert_eq(pdf.droplevel(-1), psdf.droplevel(-1)) |
| self.assert_eq(pdf.droplevel(100.0, axis=1), psdf.droplevel(100.0, axis=1)) |
| self.assert_eq(pdf.droplevel(0, axis=1), psdf.droplevel(0, axis=1)) |
| |
| def test_drop_duplicates(self): |
| pdf = pd.DataFrame( |
| {"a": [1, 2, 2, 2, 3], "b": ["a", "a", "a", "c", "d"]}, index=np.random.rand(5) |
| ) |
| psdf = ps.from_pandas(pdf) |
| |
| # inplace is False |
| for keep in ["first", "last", False]: |
| with self.subTest(keep=keep): |
| self.assert_eq( |
| pdf.drop_duplicates(keep=keep).sort_index(), |
| psdf.drop_duplicates(keep=keep).sort_index(), |
| ) |
| self.assert_eq( |
| pdf.drop_duplicates("a", keep=keep).sort_index(), |
| psdf.drop_duplicates("a", keep=keep).sort_index(), |
| ) |
| self.assert_eq( |
| pdf.drop_duplicates(["a", "b"], keep=keep).sort_index(), |
| psdf.drop_duplicates(["a", "b"], keep=keep).sort_index(), |
| ) |
| self.assert_eq( |
| pdf.set_index("a", append=True).drop_duplicates(keep=keep).sort_index(), |
| psdf.set_index("a", append=True).drop_duplicates(keep=keep).sort_index(), |
| ) |
| self.assert_eq( |
| pdf.set_index("a", append=True).drop_duplicates("b", keep=keep).sort_index(), |
| psdf.set_index("a", append=True).drop_duplicates("b", keep=keep).sort_index(), |
| ) |
| |
| columns = pd.MultiIndex.from_tuples([("x", "a"), ("y", "b")]) |
| pdf.columns = columns |
| psdf.columns = columns |
| |
| # inplace is False |
| for keep in ["first", "last", False]: |
| with self.subTest("multi-index columns", keep=keep): |
| self.assert_eq( |
| pdf.drop_duplicates(keep=keep).sort_index(), |
| psdf.drop_duplicates(keep=keep).sort_index(), |
| ) |
| self.assert_eq( |
| pdf.drop_duplicates(("x", "a"), keep=keep).sort_index(), |
| psdf.drop_duplicates(("x", "a"), keep=keep).sort_index(), |
| ) |
| self.assert_eq( |
| pdf.drop_duplicates([("x", "a"), ("y", "b")], keep=keep).sort_index(), |
| psdf.drop_duplicates([("x", "a"), ("y", "b")], keep=keep).sort_index(), |
| ) |
| self.assert_eq( |
| pdf.drop_duplicates( |
| [("x", "a"), ("y", "b")], keep=keep, ignore_index=True |
| ).sort_index(), |
| psdf.drop_duplicates( |
| [("x", "a"), ("y", "b")], keep=keep, ignore_index=True |
| ).sort_index(), |
| ) |
| |
| # inplace is True |
| subset_list = [None, "a", ["a", "b"]] |
| for subset in subset_list: |
| pdf = pd.DataFrame( |
| {"a": [1, 2, 2, 2, 3], "b": ["a", "a", "a", "c", "d"]}, index=np.random.rand(5) |
| ) |
| psdf = ps.from_pandas(pdf) |
| pser = pdf.a |
| psser = psdf.a |
| pdf.drop_duplicates(subset=subset, inplace=True) |
| psdf.drop_duplicates(subset=subset, inplace=True) |
| self.assert_eq(psdf.sort_index(), pdf.sort_index()) |
| self.assert_eq(psser.sort_index(), pser.sort_index()) |
| |
| # multi-index columns, inplace is True |
| subset_list = [None, ("x", "a"), [("x", "a"), ("y", "b")]] |
| for subset in subset_list: |
| pdf = pd.DataFrame( |
| {"a": [1, 2, 2, 2, 3], "b": ["a", "a", "a", "c", "d"]}, index=np.random.rand(5) |
| ) |
| psdf = ps.from_pandas(pdf) |
| columns = pd.MultiIndex.from_tuples([("x", "a"), ("y", "b")]) |
| pdf.columns = columns |
| psdf.columns = columns |
| pser = pdf[("x", "a")] |
| psser = psdf[("x", "a")] |
| pdf.drop_duplicates(subset=subset, inplace=True) |
| pdf.drop_duplicates(subset=subset, inplace=True, ignore_index=True) |
| psdf.drop_duplicates(subset=subset, inplace=True) |
| psdf.drop_duplicates(subset=subset, inplace=True, ignore_index=True) |
| self.assert_eq(psdf.sort_index(), pdf.sort_index()) |
| self.assert_eq(psser.sort_index(), pser.sort_index()) |
| |
| # non-string names |
| pdf = pd.DataFrame( |
| {10: [1, 2, 2, 2, 3], 20: ["a", "a", "a", "c", "d"]}, index=np.random.rand(5) |
| ) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq( |
| pdf.drop_duplicates(10, keep=keep).sort_index(), |
| psdf.drop_duplicates(10, keep=keep).sort_index(), |
| ) |
| self.assert_eq( |
| pdf.drop_duplicates([10, 20], keep=keep).sort_index(), |
| psdf.drop_duplicates([10, 20], keep=keep).sort_index(), |
| ) |
| |
| def test_duplicated(self): |
| pdf = pd.DataFrame( |
| {"a": [1, 1, 2, 3], "b": [1, 1, 1, 4], "c": [1, 1, 1, 5]}, index=np.random.rand(4) |
| ) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(pdf.duplicated().sort_index(), psdf.duplicated().sort_index()) |
| self.assert_eq( |
| pdf.duplicated(keep="last").sort_index(), |
| psdf.duplicated(keep="last").sort_index(), |
| ) |
| self.assert_eq( |
| pdf.duplicated(keep=False).sort_index(), |
| psdf.duplicated(keep=False).sort_index(), |
| ) |
| self.assert_eq( |
| pdf.duplicated(subset="b").sort_index(), |
| psdf.duplicated(subset="b").sort_index(), |
| ) |
| self.assert_eq( |
| pdf.duplicated(subset=["b"]).sort_index(), |
| psdf.duplicated(subset=["b"]).sort_index(), |
| ) |
| with self.assertRaisesRegex(ValueError, "'keep' only supports 'first', 'last' and False"): |
| psdf.duplicated(keep="false") |
| with self.assertRaisesRegex(KeyError, "'d'"): |
| psdf.duplicated(subset=["d"]) |
| |
| pdf.index.name = "x" |
| psdf.index.name = "x" |
| self.assert_eq(pdf.duplicated().sort_index(), psdf.duplicated().sort_index()) |
| |
| # multi-index |
| self.assert_eq( |
| pdf.set_index("a", append=True).duplicated().sort_index(), |
| psdf.set_index("a", append=True).duplicated().sort_index(), |
| ) |
| self.assert_eq( |
| pdf.set_index("a", append=True).duplicated(keep=False).sort_index(), |
| psdf.set_index("a", append=True).duplicated(keep=False).sort_index(), |
| ) |
| self.assert_eq( |
| pdf.set_index("a", append=True).duplicated(subset=["b"]).sort_index(), |
| psdf.set_index("a", append=True).duplicated(subset=["b"]).sort_index(), |
| ) |
| |
| # mutli-index columns |
| columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "c")]) |
| pdf.columns = columns |
| psdf.columns = columns |
| self.assert_eq(pdf.duplicated().sort_index(), psdf.duplicated().sort_index()) |
| self.assert_eq( |
| pdf.duplicated(subset=("x", "b")).sort_index(), |
| psdf.duplicated(subset=("x", "b")).sort_index(), |
| ) |
| self.assert_eq( |
| pdf.duplicated(subset=[("x", "b")]).sort_index(), |
| psdf.duplicated(subset=[("x", "b")]).sort_index(), |
| ) |
| |
| # non-string names |
| pdf = pd.DataFrame( |
| {10: [1, 1, 2, 3], 20: [1, 1, 1, 4], 30: [1, 1, 1, 5]}, index=np.random.rand(4) |
| ) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(pdf.duplicated().sort_index(), psdf.duplicated().sort_index()) |
| self.assert_eq( |
| pdf.duplicated(subset=10).sort_index(), |
| psdf.duplicated(subset=10).sort_index(), |
| ) |
| |
| def test_filter(self): |
| pdf = pd.DataFrame( |
| { |
| "aa": ["aa", "bd", "bc", "ab", "ce"], |
| "ba": [1, 2, 3, 4, 5], |
| "cb": [1.0, 2.0, 3.0, 4.0, 5.0], |
| "db": [1.0, np.nan, 3.0, np.nan, 5.0], |
| } |
| ) |
| pdf = pdf.set_index("aa") |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq( |
| psdf.filter(items=["ab", "aa"], axis=0).sort_index(), |
| pdf.filter(items=["ab", "aa"], axis=0).sort_index(), |
| ) |
| |
| with option_context("compute.isin_limit", 0): |
| self.assert_eq( |
| psdf.filter(items=["ab", "aa"], axis=0).sort_index(), |
| pdf.filter(items=["ab", "aa"], axis=0).sort_index(), |
| ) |
| |
| self.assert_eq( |
| psdf.filter(items=["ba", "db"], axis=1).sort_index(), |
| pdf.filter(items=["ba", "db"], axis=1).sort_index(), |
| ) |
| |
| self.assert_eq(psdf.filter(like="b", axis="index"), pdf.filter(like="b", axis="index")) |
| self.assert_eq(psdf.filter(like="c", axis="columns"), pdf.filter(like="c", axis="columns")) |
| |
| self.assert_eq( |
| psdf.filter(regex="b.*", axis="index"), pdf.filter(regex="b.*", axis="index") |
| ) |
| self.assert_eq( |
| psdf.filter(regex="b.*", axis="columns"), pdf.filter(regex="b.*", axis="columns") |
| ) |
| |
| pdf = pdf.set_index("ba", append=True) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq( |
| psdf.filter(items=[("aa", 1), ("bd", 2)], axis=0).sort_index(), |
| pdf.filter(items=[("aa", 1), ("bd", 2)], axis=0).sort_index(), |
| ) |
| |
| with self.assertRaisesRegex(TypeError, "Unsupported type list"): |
| psdf.filter(items=[["aa", 1], ("bd", 2)], axis=0) |
| |
| with self.assertRaisesRegex(ValueError, "The item should not be empty."): |
| psdf.filter(items=[(), ("bd", 2)], axis=0) |
| |
| self.assert_eq(psdf.filter(like="b", axis=0), pdf.filter(like="b", axis=0)) |
| |
| self.assert_eq(psdf.filter(regex="b.*", axis=0), pdf.filter(regex="b.*", axis=0)) |
| |
| with self.assertRaisesRegex(ValueError, "items should be a list-like object"): |
| psdf.filter(items="b") |
| |
| with self.assertRaisesRegex(ValueError, "No axis named"): |
| psdf.filter(regex="b.*", axis=123) |
| |
| with self.assertRaisesRegex(TypeError, "Must pass either `items`, `like`"): |
| psdf.filter() |
| |
| with self.assertRaisesRegex(TypeError, "mutually exclusive"): |
| psdf.filter(regex="b.*", like="aaa") |
| |
| # multi-index columns |
| pdf = pd.DataFrame( |
| { |
| ("x", "aa"): ["aa", "ab", "bc", "bd", "ce"], |
| ("x", "ba"): [1, 2, 3, 4, 5], |
| ("y", "cb"): [1.0, 2.0, 3.0, 4.0, 5.0], |
| ("z", "db"): [1.0, np.nan, 3.0, np.nan, 5.0], |
| } |
| ) |
| pdf = pdf.set_index(("x", "aa")) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq( |
| psdf.filter(items=["ab", "aa"], axis=0).sort_index(), |
| pdf.filter(items=["ab", "aa"], axis=0).sort_index(), |
| ) |
| self.assert_eq( |
| psdf.filter(items=[("x", "ba"), ("z", "db")], axis=1).sort_index(), |
| pdf.filter(items=[("x", "ba"), ("z", "db")], axis=1).sort_index(), |
| ) |
| |
| self.assert_eq(psdf.filter(like="b", axis="index"), pdf.filter(like="b", axis="index")) |
| self.assert_eq(psdf.filter(like="c", axis="columns"), pdf.filter(like="c", axis="columns")) |
| |
| self.assert_eq( |
| psdf.filter(regex="b.*", axis="index"), pdf.filter(regex="b.*", axis="index") |
| ) |
| self.assert_eq( |
| psdf.filter(regex="b.*", axis="columns"), pdf.filter(regex="b.*", axis="columns") |
| ) |
| |
| def test_last(self): |
| index = pd.date_range("2018-04-09", periods=4, freq="2D") |
| pdf = pd.DataFrame([1, 2, 3, 4], index=index) |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq(pdf.last("1D"), psdf.last("1D")) |
| self.assert_eq(pdf.last(DateOffset(days=1)), psdf.last(DateOffset(days=1))) |
| with self.assertRaisesRegex(TypeError, "'last' only supports a DatetimeIndex"): |
| ps.DataFrame([1, 2, 3, 4]).last("1D") |
| |
| def test_first(self): |
| index = pd.date_range("2018-04-09", periods=4, freq="2D") |
| pdf = pd.DataFrame([1, 2, 3, 4], index=index) |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq(pdf.first("1D"), psdf.first("1D")) |
| self.assert_eq(pdf.first(DateOffset(days=1)), psdf.first(DateOffset(days=1))) |
| with self.assertRaisesRegex(TypeError, "'first' only supports a DatetimeIndex"): |
| ps.DataFrame([1, 2, 3, 4]).first("1D") |
| |
| 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")) |
| pdf = pd.DataFrame({"x1": ["a", "b", "c", "d"], "x2": ["a", "b", "c", "d"]}, index=pidx) |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq(pdf.swaplevel(), psdf.swaplevel()) |
| self.assert_eq(pdf.swaplevel(0, 1), psdf.swaplevel(0, 1)) |
| self.assert_eq(pdf.swaplevel(1, 1), psdf.swaplevel(1, 1)) |
| self.assert_eq(pdf.swaplevel("number", "color"), psdf.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")) |
| pdf = pd.DataFrame({"x1": ["a", "b", "c", "d"], "x2": ["a", "b", "c", "d"]}, index=pidx) |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq(pdf.swaplevel(), psdf.swaplevel()) |
| self.assert_eq(pdf.swaplevel(0, 1), psdf.swaplevel(0, 1)) |
| self.assert_eq(pdf.swaplevel(0, 2), psdf.swaplevel(0, 2)) |
| self.assert_eq(pdf.swaplevel(1, 2), psdf.swaplevel(1, 2)) |
| self.assert_eq(pdf.swaplevel(1, 1), psdf.swaplevel(1, 1)) |
| self.assert_eq(pdf.swaplevel(-1, -2), psdf.swaplevel(-1, -2)) |
| self.assert_eq(pdf.swaplevel("number", "color"), psdf.swaplevel("number", "color")) |
| self.assert_eq(pdf.swaplevel("number", "size"), psdf.swaplevel("number", "size")) |
| self.assert_eq(pdf.swaplevel("color", "size"), psdf.swaplevel("color", "size")) |
| self.assert_eq( |
| pdf.swaplevel("color", "size", axis="index"), |
| psdf.swaplevel("color", "size", axis="index"), |
| ) |
| self.assert_eq( |
| pdf.swaplevel("color", "size", axis=0), psdf.swaplevel("color", "size", axis=0) |
| ) |
| |
| pdf = pd.DataFrame( |
| { |
| "x1": ["a", "b", "c", "d"], |
| "x2": ["a", "b", "c", "d"], |
| "x3": ["a", "b", "c", "d"], |
| "x4": ["a", "b", "c", "d"], |
| } |
| ) |
| pidx = pd.MultiIndex.from_arrays(arrays, names=("number", "color", "size")) |
| pdf.columns = pidx |
| psdf = ps.from_pandas(pdf) |
| self.assert_eq(pdf.swaplevel(axis=1), psdf.swaplevel(axis=1)) |
| self.assert_eq(pdf.swaplevel(0, 1, axis=1), psdf.swaplevel(0, 1, axis=1)) |
| self.assert_eq(pdf.swaplevel(0, 2, axis=1), psdf.swaplevel(0, 2, axis=1)) |
| self.assert_eq(pdf.swaplevel(1, 2, axis=1), psdf.swaplevel(1, 2, axis=1)) |
| self.assert_eq(pdf.swaplevel(1, 1, axis=1), psdf.swaplevel(1, 1, axis=1)) |
| self.assert_eq(pdf.swaplevel(-1, -2, axis=1), psdf.swaplevel(-1, -2, axis=1)) |
| self.assert_eq( |
| pdf.swaplevel("number", "color", axis=1), psdf.swaplevel("number", "color", axis=1) |
| ) |
| self.assert_eq( |
| pdf.swaplevel("number", "size", axis=1), psdf.swaplevel("number", "size", axis=1) |
| ) |
| self.assert_eq( |
| pdf.swaplevel("color", "size", axis=1), psdf.swaplevel("color", "size", axis=1) |
| ) |
| self.assert_eq( |
| pdf.swaplevel("color", "size", axis="columns"), |
| psdf.swaplevel("color", "size", axis="columns"), |
| ) |
| |
| # Error conditions |
| self.assertRaises(AssertionError, lambda: ps.DataFrame([1, 2]).swaplevel()) |
| self.assertRaises(IndexError, lambda: psdf.swaplevel(0, 9, axis=1)) |
| self.assertRaises(KeyError, lambda: psdf.swaplevel("not_number", "color", axis=1)) |
| self.assertRaises(ValueError, lambda: psdf.swaplevel(axis=2)) |
| |
| def test_swapaxes(self): |
| pdf = pd.DataFrame( |
| [[1, 2, 3], [4, 5, 6], [7, 8, 9]], index=["x", "y", "z"], columns=["a", "b", "c"] |
| ) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(psdf.swapaxes(0, 1), pdf.swapaxes(0, 1)) |
| self.assert_eq(psdf.swapaxes(1, 0), pdf.swapaxes(1, 0)) |
| self.assert_eq(psdf.swapaxes("index", "columns"), pdf.swapaxes("index", "columns")) |
| self.assert_eq(psdf.swapaxes("columns", "index"), pdf.swapaxes("columns", "index")) |
| self.assert_eq((psdf + 1).swapaxes(0, 1), (pdf + 1).swapaxes(0, 1)) |
| |
| self.assertRaises(AssertionError, lambda: psdf.swapaxes(0, 1, copy=False)) |
| self.assertRaises(ValueError, lambda: psdf.swapaxes(0, -1)) |
| |
| def test_isin(self): |
| pdf = pd.DataFrame( |
| { |
| "a": [4, 2, 3, 4, 8, 6], |
| "b": [1, 2, 9, 4, 2, 4], |
| "c": ["one", "three", "six", "seven", "one", "5"], |
| }, |
| index=np.random.rand(6), |
| ) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(psdf.isin([4, "six"]), pdf.isin([4, "six"])) |
| # Seems like pandas has a bug when passing `np.array` as parameter |
| self.assert_eq(psdf.isin(np.array([4, "six"])), pdf.isin([4, "six"])) |
| self.assert_eq( |
| psdf.isin({"a": [2, 8], "c": ["three", "one"]}), |
| pdf.isin({"a": [2, 8], "c": ["three", "one"]}), |
| ) |
| self.assert_eq( |
| psdf.isin({"a": np.array([2, 8]), "c": ["three", "one"]}), |
| pdf.isin({"a": np.array([2, 8]), "c": ["three", "one"]}), |
| ) |
| |
| msg = "'DataFrame' object has no attribute {'e'}" |
| with self.assertRaisesRegex(AttributeError, msg): |
| psdf.isin({"e": [5, 7], "a": [1, 6]}) |
| |
| msg = "DataFrame and Series are not supported" |
| with self.assertRaisesRegex(NotImplementedError, msg): |
| psdf.isin(pdf) |
| |
| msg = "Values should be iterable, Series, DataFrame or dict." |
| with self.assertRaisesRegex(TypeError, msg): |
| psdf.isin(1) |
| |
| pdf = pd.DataFrame( |
| { |
| "a": [4, 2, 3, 4, 8, 6], |
| "b": [1, None, 9, 4, None, 4], |
| "c": [None, 5, None, 3, 2, 1], |
| }, |
| ) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(psdf.isin([4, 3, 1, 1, None]), pdf.isin([4, 3, 1, 1, None])) |
| |
| self.assert_eq(psdf.isin({"b": [4, 3, 1, 1, None]}), pdf.isin({"b": [4, 3, 1, 1, None]})) |
| |
| def test_sample(self): |
| psdf = ps.DataFrame({"A": [0, 2, 4]}, index=["x", "y", "z"]) |
| |
| # Make sure the tests run, but we can't check the result because they are non-deterministic. |
| psdf.sample(frac=0.1) |
| psdf.sample(frac=0.2, replace=True) |
| psdf.sample(frac=0.2, random_state=5) |
| psdf["A"].sample(frac=0.2) |
| psdf["A"].sample(frac=0.2, replace=True) |
| psdf["A"].sample(frac=0.2, random_state=5) |
| |
| self.assert_eq(psdf.sample(frac=0.1, ignore_index=True).index.dtype, np.int64) |
| self.assert_eq(psdf.sample(frac=0.2, replace=True, ignore_index=True).index.dtype, np.int64) |
| self.assert_eq( |
| psdf.sample(frac=0.2, random_state=5, ignore_index=True).index.dtype, np.int64 |
| ) |
| self.assert_eq(psdf["A"].sample(frac=0.2, ignore_index=True).index.dtype, np.int64) |
| self.assert_eq( |
| psdf["A"].sample(frac=0.2, replace=True, ignore_index=True).index.dtype, np.int64 |
| ) |
| self.assert_eq( |
| psdf["A"].sample(frac=0.2, random_state=5, ignore_index=True).index.dtype, np.int64 |
| ) |
| |
| with self.assertRaises(ValueError): |
| psdf.sample() |
| with self.assertRaises(NotImplementedError): |
| psdf.sample(n=1) |
| |
| |
| class FrameReidexingTests( |
| FrameReindexingMixin, |
| PandasOnSparkTestCase, |
| SQLTestUtils, |
| ): |
| pass |
| |
| |
| if __name__ == "__main__": |
| from pyspark.pandas.tests.frame.test_reindexing 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) |
