From 45553c5b3090230e23c8deca221354dad35c52de Mon Sep 17 00:00:00 2001
From: karajan1001 <mishanyo1001@gmail.com>
Date: Sun, 31 Oct 2021 21:20:43 +0800
Subject: [PATCH] TabularData: add a new how `arg` to dropna(#6892)

Borrowed this idea from pandas api. Current dropna methods will drop a
column or row if any data is missing. But it is also quite useful to
drop those columns or rows only if all of its elements are missing. This
is mostly used for pruning a table.
---
 dvc/compare.py                  | 29 ++++++++++++++++++++----
 tests/unit/test_tabular_data.py | 39 +++++++++++++++++++++++----------
 2 files changed, 52 insertions(+), 16 deletions(-)

diff --git a/dvc/compare.py b/dvc/compare.py
index e9d742d3b..9af04f35d 100644
--- a/dvc/compare.py
+++ b/dvc/compare.py
@@ -182,21 +182,42 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
             {k: self._columns[k][i] for k in keys} for i in range(len(self))
         ]
 
-    def dropna(self, axis: str = "rows"):
+    def dropna(self, axis: str = "rows", how="any"):
         if axis not in ["rows", "cols"]:
             raise ValueError(
                 f"Invalid 'axis' value {axis}."
                 "Choose one of ['rows', 'cols']"
             )
+<<<<<<< HEAD
         to_drop: Set = set()
+=======
+        if how not in ["any", "all"]:
+            raise ValueError(
+                f"Invalid 'how' value {how}." "Choose one of ['any', 'all']"
+            )
+
+        match_line: Set = set()
+        match = True
+        if how == "all":
+            match = False
+
+>>>>>>> 8ad4b8aa (TabularData: add a new how `arg` to dropna(#6892))
         for n_row, row in enumerate(self):
             for n_col, col in enumerate(row):
-                if col == self._fill_value:
+                if (col == self._fill_value) is match:
                     if axis == "rows":
-                        to_drop.add(n_row)
+                        match_line.add(n_row)
                         break
                     else:
-                        to_drop.add(self.keys()[n_col])
+                        match_line.add(self.keys()[n_col])
+
+        to_drop = match_line
+        if how == "all":
+            if axis == "rows":
+                to_drop = set(range(len(self)))
+            else:
+                to_drop = set(self.keys())
+            to_drop -= match_line
 
         if axis == "rows":
             for name in self.keys():
diff --git a/tests/unit/test_tabular_data.py b/tests/unit/test_tabular_data.py
index 43c9840b5..a8439d281 100644
--- a/tests/unit/test_tabular_data.py
+++ b/tests/unit/test_tabular_data.py
@@ -180,28 +180,43 @@ def test_row_from_dict():
 
 
 @pytest.mark.parametrize(
-    "axis,expected",
+    "axis,how,data,expected",
     [
         (
             "rows",
+            "any",
+            [["foo"], ["foo", "bar"], ["foo", "bar", "foobar"]],
             [
                 ["foo", "bar", "foobar"],
             ],
         ),
-        ("cols", [["foo"], ["foo"], ["foo"]]),
+        (
+            "rows",
+            "all",
+            [["foo"], ["foo", "bar"], ["", "", ""]],
+            [
+                ["foo", "", ""],
+                ["foo", "bar", ""],
+            ],
+        ),
+        (
+            "cols",
+            "any",
+            [["foo"], ["foo", "bar"], ["foo", "bar", "foobar"]],
+            [["foo"], ["foo"], ["foo"]],
+        ),
+        (
+            "cols",
+            "all",
+            [["foo"], ["foo", "bar"], ["", "", ""]],
+            [["foo", ""], ["foo", "bar"], ["", ""]],
+        ),
     ],
 )
-def test_dropna(axis, expected):
+def test_dropna(axis, how, data, expected):
     td = TabularData(["col-1", "col-2", "col-3"])
-    td.extend([["foo"], ["foo", "bar"], ["foo", "bar", "foobar"]])
-    assert list(td) == [
-        ["foo", "", ""],
-        ["foo", "bar", ""],
-        ["foo", "bar", "foobar"],
-    ]
-
-    td.dropna(axis)
-
+    td.extend(data)
+    td.dropna(axis, how)
     assert list(td) == expected