From 255e6a1aecb2fc34e89c0721b36fdf0db4789bdf Mon Sep 17 00:00:00 2001 From: David de la Iglesia Castro Date: Mon, 3 Jan 2022 20:33:37 +0100 Subject: [PATCH] `TabularData`: Support protected`columns. TabularData: Add subset arg to dropna and drop_duplicates. Closes #7083 --- dvc/compare.py | 49 ++++++++++++--- tests/unit/test_tabular_data.py | 108 +++++++++++++++++++++++--------- 2 files changed, 121 insertions(+), 36 deletions(-) diff --git a/dvc/compare.py b/dvc/compare.py index f9c3ffa12..9037cd637 100644 --- a/dvc/compare.py +++ b/dvc/compare.py @@ -11,6 +11,7 @@ from typing import ( List, Mapping, MutableSequence, + Optional, Sequence, Set, Tuple, @@ -38,11 +39,21 @@ class TabularData(MutableSequence[Sequence["CellT"]]): self._columns: Dict[str, Column] = {name: Column() for name in columns} self._keys: List[str] = list(columns) self._fill_value = fill_value + self._protected: Set[str] = set() @property def columns(self) -> List[Column]: return list(map(self.column, self.keys())) + def is_protected(self, col_name) -> bool: + return col_name in self._protected + + def protect(self, *col_names: str): + self._protected.update(col_names) + + def unprotect(self, *col_names: str): + self._protected = self._protected.difference(col_names) + def column(self, name: str) -> Column: return self._columns[name] @@ -126,9 +137,13 @@ class TabularData(MutableSequence[Sequence["CellT"]]): return len(self.columns), len(self) def drop(self, *col_names: str) -> None: - for col_name in col_names: - self._keys.remove(col_name) - self._columns.pop(col_name) + to_remove = set() + for col in col_names: + if not self.is_protected(col): + to_remove.add(col) + for col in to_remove: + self._keys.remove(col) + self._columns.pop(col) def rename(self, from_col_name: str, to_col_name: str) -> None: self._columns[to_col_name] = self._columns.pop(from_col_name) @@ -199,7 +214,9 @@ 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", how="any"): + def dropna( + self, axis: str = "rows", how="any", subset: Optional[List] = None + ): if axis not in ["rows", "cols"]: raise ValueError( f"Invalid 'axis' value {axis}." @@ -211,13 +228,15 @@ class TabularData(MutableSequence[Sequence["CellT"]]): ) match_line: Set = set() - match = True + match_any = True if how == "all": - match = False + match_any = False for n_row, row in enumerate(self): for n_col, col in enumerate(row): - if (col == self._fill_value) is match: + if subset and self.keys()[n_col] not in subset: + continue + if (col == self._fill_value) is match_any: if axis == "rows": match_line.add(n_row) break @@ -244,7 +263,12 @@ class TabularData(MutableSequence[Sequence["CellT"]]): else: self.drop(*to_drop) - def drop_duplicates(self, axis: str = "rows", ignore_empty: bool = True): + def drop_duplicates( + self, + axis: str = "rows", + subset: Optional[List] = None, + ignore_empty: bool = True, + ): if axis not in ["rows", "cols"]: raise ValueError( f"Invalid 'axis' value {axis}." @@ -254,6 +278,8 @@ class TabularData(MutableSequence[Sequence["CellT"]]): if axis == "cols": cols_to_drop: List[str] = [] for n_col, col in enumerate(self.columns): + if subset and self.keys()[n_col] not in subset: + continue # Cast to str because Text is not hashable error unique_vals = {str(x) for x in col} if ignore_empty and self._fill_value in unique_vals: @@ -266,6 +292,13 @@ class TabularData(MutableSequence[Sequence["CellT"]]): unique_rows = [] rows_to_drop: List[int] = [] for n_row, row in enumerate(self): + if subset: + row = [ + col + for n_col, col in enumerate(row) + if self.keys()[n_col] in subset + ] + tuple_row = tuple(row) if tuple_row in unique_rows: rows_to_drop.append(n_row) diff --git a/tests/unit/test_tabular_data.py b/tests/unit/test_tabular_data.py index 328bce71f..7b6b0f409 100644 --- a/tests/unit/test_tabular_data.py +++ b/tests/unit/test_tabular_data.py @@ -144,12 +144,28 @@ def test_fill_value(): def test_drop(): - td = TabularData(["col1", "col2", "col3"]) - td.append(["foo", "bar", "baz"]) - assert list(td) == [["foo", "bar", "baz"]] + td = TabularData(["col1", "col2", "col3", "other"]) + td.append(["foo", "bar", "baz", "other_val"]) + assert list(td) == [["foo", "bar", "baz", "other_val"]] td.drop("col2") - assert td.keys() == ["col1", "col3"] - assert list(td) == [["foo", "baz"]] + assert td.keys() == ["col1", "col3", "other"] + assert list(td) == [["foo", "baz", "other_val"]] + + +def test_protected(): + td = TabularData(["col1", "col2", "col3", "other"]) + td.append(["foo", "bar", "baz", "other_val"]) + td.protect("col1", "col2") + + td.drop("col1", "col2", "col3", "other") + assert td.keys() == ["col1", "col2"] + assert list(td) == [["foo", "bar"]] + + td.unprotect("col2") + + td.drop("col1", "col2") + assert td.keys() == ["col1"] + assert list(td) == [["foo"]] def test_row_from_dict(): @@ -220,6 +236,20 @@ def test_dropna(axis, how, data, expected): assert list(td) == expected +@pytest.mark.parametrize( + "axis,expected", + [ + ("cols", [["foo", ""], ["foo", ""], ["foo", "foobar"]]), + ("rows", [["foo", "bar", ""], ["foo", "bar", "foobar"]]), + ], +) +def test_dropna_subset(axis, expected): + td = TabularData(["col-1", "col-2", "col-3"]) + td.extend([["foo"], ["foo", "bar"], ["foo", "bar", "foobar"]]) + td.dropna(axis, subset=["col-1", "col-2"]) + assert list(td) == expected + + @pytest.mark.parametrize( "axis,expected,ignore_empty", [ @@ -263,29 +293,6 @@ def test_drop_duplicates(axis, expected, ignore_empty): assert list(td) == expected -def test_drop_duplicates_ignore_empty(): - td = TabularData(["col-1", "col-2", "col-3"], fill_value="-") - td.extend( - [["foo"], ["foo", "foo"], ["foo", "foo"], ["foo", "bar", "foobar"]] - ) - - assert list(td) == [ - ["foo", "-", "-"], - ["foo", "foo", "-"], - ["foo", "foo", "-"], - ["foo", "bar", "foobar"], - ] - - td.drop_duplicates("cols", ignore_empty=False) - - assert list(td) == [ - ["-", "-"], - ["foo", "-"], - ["foo", "-"], - ["bar", "foobar"], - ] - - def test_drop_duplicates_rich_text(): from dvc.ui import ui @@ -312,6 +319,51 @@ def test_drop_duplicates_rich_text(): assert list(td) == [["-"], ["foo"], ["foo"], ["bar"]] +@pytest.mark.parametrize( + "axis,subset,expected", + [ + ( + "rows", + ["col-1"], + [["foo", "foo", "foo", "bar"]], + ), + ( + "rows", + ["col-1", "col-3"], + [ + ["foo", "foo", "foo", "bar"], + ["foo", "bar", "foobar", "bar"], + ], + ), + ( + "cols", + ["col-1", "col-3"], + [ + ["foo", "foo", "bar"], + ["bar", "foo", "bar"], + ["bar", "foobar", "bar"], + ], + ), + ], +) +def test_drop_duplicates_subset(axis, subset, expected): + td = TabularData(["col-1", "col-2", "col-3", "col-4"]) + td.extend( + [ + ["foo", "foo", "foo", "bar"], + ["foo", "bar", "foo", "bar"], + ["foo", "bar", "foobar", "bar"], + ] + ) + assert list(td) == [ + ["foo", "foo", "foo", "bar"], + ["foo", "bar", "foo", "bar"], + ["foo", "bar", "foobar", "bar"], + ] + td.drop_duplicates(axis, subset=subset) + assert list(td) == expected + + def test_dropna_invalid_axis(): td = TabularData(["col-1", "col-2", "col-3"])