Closes #3020 dataframe.dropna (#3101)

Co-authored-by: Amanda Potts <[email protected]>
Bears-R-Us · Apr 19, 2024 · 8ac2645 · 8ac2645
1 parent 8c1c57e
commit 8ac2645
Show file tree

Hide file tree

Showing 4 changed files with 718 additions and 3 deletions.
diff --git a/PROTO_tests/tests/dataframe_test.py b/PROTO_tests/tests/dataframe_test.py
@@ -1002,6 +1002,124 @@ def test_multi_col_merge(self):
                     # assert_frame_equal(sorted_ak.to_pandas()[sorted_column_names],
                     # sorted_pd[sorted_column_names])
 
+    def test_isna_notna(self):
+        df = ak.DataFrame(
+            {
+                "A": [np.nan, 2, 2, 3],
+                "B": [3, np.nan, 5, 0],
+                "C": [1, np.nan, 2, np.nan],
+                "D": ["a", "b", "c", ""],
+            }
+        )
+        assert_frame_equal(df.isna().to_pandas(), df.to_pandas().isna())
+        assert_frame_equal(df.notna().to_pandas(), df.to_pandas().notna())
+
+    def test_any_all(self):
+        df1 = ak.DataFrame(
+            {
+                "A": [True, True, True, True],
+                "B": [True, True, True, False],
+                "C": [True, False, True, False],
+                "D": [False, False, False, False],
+                "E": [0, 1, 2, 3],
+                "F": ["a", "b", "c", ""],
+            }
+        )
+
+        df2 = ak.DataFrame(
+            {
+                "A": [True, True, True, True],
+                "B": [True, True, True, True],
+            }
+        )
+
+        df3 = ak.DataFrame(
+            {
+                "A": [False, False],
+                "B": [False, False],
+            }
+        )
+
+        df4 = ak.DataFrame({"A": [1, 2, 3], "B": ["a", "b", "c"]})
+
+        df5 = ak.DataFrame()
+
+        for df in [df1, df2, df3, df4, df5]:
+            for axis in [0, 1, "index", "columns"]:
+                # There's a bug in assert_series_equal where two empty series will not register as equal.
+                if df.to_pandas().any(axis=axis, bool_only=True).empty:
+                    assert df.any(axis=axis).to_pandas().empty is True
+                else:
+                    assert_series_equal(
+                        df.any(axis=axis).to_pandas(), df.to_pandas().any(axis=axis, bool_only=True)
+                    )
+                if df.to_pandas().all(axis=axis, bool_only=True).empty:
+                    assert df.all(axis=axis).to_pandas().empty is True
+                else:
+                    assert_series_equal(
+                        df.all(axis=axis).to_pandas(), df.to_pandas().all(axis=axis, bool_only=True)
+                    )
+            # Test is axis=None
+            assert df.any(axis=None) == df.to_pandas().any(axis=None, bool_only=True)
+            assert df.all(axis=None) == df.to_pandas().all(axis=None, bool_only=True)
+
+    def test_dropna(self):
+        df1 = ak.DataFrame(
+            {
+                "A": [True, True, True, True],
+                "B": [1, np.nan, 2, np.nan],
+                "C": [1, 2, 3, np.nan],
+                "D": [False, False, False, False],
+                "E": [1, 2, 3, 4],
+                "F": ["a", "b", "c", "d"],
+                "G": [1, 2, 3, 4],
+            }
+        )
+
+        df2 = ak.DataFrame(
+            {
+                "A": [True, True, True, True],
+                "B": [True, True, True, True],
+            }
+        )
+
+        df3 = ak.DataFrame(
+            {
+                "A": [False, False],
+                "B": [False, False],
+            }
+        )
+
+        df4 = ak.DataFrame({"A": [1, 2, 3], "B": ["a", "b", "c"]})
+
+        df5 = ak.DataFrame()
+
+        for df in [df1, df2, df3, df4, df5]:
+            for axis in [0, 1, "index", "columns"]:
+                for how in ["any", "all"]:
+                    for ignore_index in [True, False]:
+                        assert_frame_equal(
+                            df.dropna(axis=axis, how=how, ignore_index=ignore_index).to_pandas(
+                                retain_index=True
+                            ),
+                            df.to_pandas(retain_index=True).dropna(
+                                axis=axis, how=how, ignore_index=ignore_index
+                            ),
+                        )
+
+                for thresh in [0, 1, 2, 3, 4, 5]:
+                    if df.to_pandas(retain_index=True).dropna(axis=axis, thresh=thresh).empty:
+                        assert (
+                            df.dropna(axis=axis, thresh=thresh).to_pandas(retain_index=True).empty
+                            == True
+                        )
+
+                    else:
+                        assert_frame_equal(
+                            df.dropna(axis=axis, thresh=thresh).to_pandas(retain_index=True),
+                            df.to_pandas(retain_index=True).dropna(axis=axis, thresh=thresh),
+                        )
+
     def test_memory_usage(self):
         dtypes = [ak.int64, ak.float64, ak.bool]
         data = dict([(str(t), ak.ones(5000, dtype=ak.int64).astype(t)) for t in dtypes])