Resolve review

jupyter-widgets · Mar 21, 2024 · 0d72d90 · 0d72d90
1 parent 19f397f
commit 0d72d90
Showing 1 changed file with 2 additions and 7 deletions.
diff --git a/ipydatagrid/datagrid.py b/ipydatagrid/datagrid.py
@@ -896,13 +896,8 @@ def _get_row_index_of_primary_key(self, value):
                 "as the primary key."
             )
 
-        # TODO Is there a better way for this?
-        row_indices = [
-            idx
-            for idx, row in self._data["data"].iterrows()
-            if all(row[key[j]] == value[j] for j in range(len(key)))
-        ]
-        return row_indices
+        df = self._data["data"]
+        return pd.RangeIndex(len(df))[(df[key] == value).all(axis="columns")].to_list()
 
     @staticmethod
     def _get_cell_value_by_numerical_index(data, column_index, row_index):