From cb6f625dd9dc5698142d2b4969ed3a766534511b Mon Sep 17 00:00:00 2001
From: David de la Iglesia Castro <daviddelaiglesiacastro@gmail.com>
Date: Wed, 27 Oct 2021 19:26:04 +0200
Subject: [PATCH] TabularData: Add `drop_duplicates` method.

This method mimics pandas `drop_duplicates` when axis=rows.

For axis=cols this method will drop columns with 0-variance (only
1 unique value). NaNs (`_fill_value`) won't be counted as unique values.

So the following column (`-` would be the `_fill_value`) will be dropped:

foo
-
foo
foo
---
 dvc/compare.py                  | 36 +++++++++++++++++++++++++++++-
 tests/unit/test_tabular_data.py | 39 +++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/dvc/compare.py b/dvc/compare.py
index 9384d5f1f..e9d742d3b 100644
--- a/dvc/compare.py
+++ b/dvc/compare.py
@@ -188,7 +188,6 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
                 f"Invalid 'axis' value {axis}."
                 "Choose one of ['rows', 'cols']"
             )
-
         to_drop: Set = set()
         for n_row, row in enumerate(self):
             for n_col, col in enumerate(row):
@@ -211,6 +210,41 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
         else:
             self.drop(*to_drop)
 
+    def drop_duplicates(self, axis: str = "rows"):
+        if axis not in ["rows", "cols"]:
+            raise ValueError(
+                f"Invalid 'axis' value {axis}."
+                "Choose one of ['rows', 'cols']"
+            )
+
+        if axis == "cols":
+            cols_to_drop: List[str] = []
+            for n_col, col in enumerate(self.columns):
+                # Cast to str because Text is not hashable error
+                unique_vals = {str(x) for x in col if x != self._fill_value}
+                if len(unique_vals) == 1:
+                    cols_to_drop.append(self.keys()[n_col])
+            self.drop(*cols_to_drop)
+
+        elif axis == "rows":
+            unique_rows = []
+            rows_to_drop: List[int] = []
+            for n_row, row in enumerate(self):
+                tuple_row = tuple(row)
+                if tuple_row in unique_rows:
+                    rows_to_drop.append(n_row)
+                else:
+                    unique_rows.append(tuple_row)
+
+            for name in self.keys():
+                self._columns[name] = Column(
+                    [
+                        x
+                        for n, x in enumerate(self._columns[name])
+                        if n not in rows_to_drop
+                    ]
+                )
+
 
 def _normalize_float(val: float, precision: int):
     return f"{val:.{precision}g}"
diff --git a/tests/unit/test_tabular_data.py b/tests/unit/test_tabular_data.py
index 9631c4362..43c9840b5 100644
--- a/tests/unit/test_tabular_data.py
+++ b/tests/unit/test_tabular_data.py
@@ -205,8 +205,47 @@ def test_dropna(axis, expected):
     assert list(td) == expected
 
 
+@pytest.mark.parametrize(
+    "axis,expected",
+    [
+        (
+            "rows",
+            [
+                ["foo", "", ""],
+                ["foo", "foo", ""],
+                ["foo", "bar", "foobar"],
+            ],
+        ),
+        ("cols", [[""], ["foo"], ["foo"], ["bar"]]),
+    ],
+)
+def test_drop_duplicates(axis, expected):
+    td = TabularData(["col-1", "col-2", "col-3"])
+    td.extend(
+        [["foo"], ["foo", "foo"], ["foo", "foo"], ["foo", "bar", "foobar"]]
+    )
+
+    assert list(td) == [
+        ["foo", "", ""],
+        ["foo", "foo", ""],
+        ["foo", "foo", ""],
+        ["foo", "bar", "foobar"],
+    ]
+
+    td.drop_duplicates(axis)
+
+    assert list(td) == expected
+
+
 def test_dropna_invalid_axis():
     td = TabularData(["col-1", "col-2", "col-3"])
 
     with pytest.raises(ValueError, match="Invalid 'axis' value foo."):
         td.dropna("foo")
+
+
+def test_drop_duplicates_invalid_axis():
+    td = TabularData(["col-1", "col-2", "col-3"])
+
+    with pytest.raises(ValueError, match="Invalid 'axis' value foo."):
+        td.drop_duplicates("foo")