Optimize Dataset.subset() lookups

AUTHORS

@@ -12,6 +12,7 @@ Here is a list of past and present much-appreciated contributors:
     Daniel Santos
     Egor Osokin
     Erik Youngren
+    Harshit Gupta
     Hugo van Kemenade
     Ian Stride
     Iuri de Silvio

src/tablib/core.py

@@ -836,24 +836,21 @@ def subset(self, rows=None, cols=None):
 
         # filter out impossible rows and columns
         rows = [row for row in rows if row in range(self.height)]
+        # Keep list membership for selectors with custom equality or hashing.
+        if all(type(row) is int for row in rows):
+            rows = set(rows)
         cols = [header for header in cols if header in self.headers]
 
         _dset = Dataset()
 
         # filtering rows and columns
         _dset.headers = list(cols)
+        col_indexes = [self.headers.index(header) for header in _dset.headers]
 
         _dset._data = []
         for row_no, row in enumerate(self._data):
-            data_row = []
-            for key in _dset.headers:
-                if key in self.headers:
-                    pos = self.headers.index(key)
-                    data_row.append(row[pos])
-                else:
-                    raise KeyError
-
             if row_no in rows:
+                data_row = [row[index] for index in col_indexes]
                 _dset.append(row=Row(data_row))
 
         return _dset

tests/test_tablib.py

@@ -610,6 +610,46 @@ def test_subset(self):
         self.assertEqual(subset._data[0].list, ['John', 90])
         self.assertEqual(subset._data[1].list, ['Thomas', 50])
 
+    def test_subset_preserves_selection_behavior(self):
+        data.headers = self.headers
+        data.append(self.john)
+        data.append(self.george)
+        data.append(self.tom)
+
+        subset = data.subset(
+            rows=(2, 0, 2, -1, 99),
+            cols=('gpa', 'missing', 'first_name', 'gpa'),
+        )
+
+        self.assertEqual(subset.headers, ['gpa', 'first_name', 'gpa'])
+        self.assertEqual(subset._data[0].list, [90, 'John', 90])
+        self.assertEqual(subset._data[1].list, [50, 'Thomas', 50])
+
+    def test_subset_preserves_no_matching_column_behavior(self):
+        data.headers = self.headers
+        data.append(self.john)
+
+        for columns in ([], ['missing']):
+            with self.subTest(columns=columns):
+                with self.assertRaisesRegex(TypeError, "'NoneType' object is not iterable"):
+                    data.subset(cols=columns)
+
+    def test_subset_preserves_unhashable_row_selection(self):
+        class FirstRow:
+            __hash__ = None
+
+            def __eq__(self, other):
+                return other == 0
+
+        data.headers = self.headers
+        data.append(self.john)
+        data.append(self.george)
+
+        subset = data.subset(rows=[1, FirstRow()], cols=['first_name'])
+        self.assertEqual(subset[0], ('John',))
+        self.assertEqual(subset[1], ('George',))
+        self.assertEqual(subset.height, 2)
+
     def test_formatters(self):
         """Confirm formatters are being triggered."""