rewrite to preserve types

Author: stevenae

pandas/core/frame.py

@@ -101,7 +101,6 @@
     is_integer_dtype,
     is_iterator,
     is_list_like,
-    is_object_dtype,
     is_scalar,
     is_sequence,
     needs_i8_conversion,
@@ -5204,17 +5203,13 @@ def lookup(self, row_labels, col_labels) -> ExtensionArray | np.ndarray:
             raise KeyError("One or more column labels was not found")
 
         sub = self.take(np.unique(cidx), axis=1)
-        if sub._is_mixed_type:
-            sub = sub.take(np.unique(ridx), axis=0)
-            ridx = sub.index.get_indexer(row_labels)
-        values = sub.values
+        sub = sub.take(np.unique(ridx), axis=0)
+        ridx = sub.index.get_indexer(row_labels)
+        values = sub.melt()["value"]
         cidx = sub.columns.get_indexer(col_labels)
-        flat_index = ridx * len(sub.columns) + cidx
+        flat_index = ridx + cidx * len(sub)
 
-        result = values.flat[flat_index]
-
-        if is_object_dtype(result):
-            result = lib.maybe_convert_objects(result)
+        result = values[flat_index]
 
         return result
 

pandas/tests/frame/indexing/test_indexing.py

@@ -1420,19 +1420,17 @@ def test_lookup_float(self, float_frame):
         cols = list(df.columns) * len(df.index)
         result = df.lookup(rows, cols)
 
-        expected = np.array([df.loc[r, c] for r, c in zip(rows, cols)])
-        tm.assert_numpy_array_equal(result, expected)
+        expected = Series([df.loc[r, c] for r, c in zip(rows, cols)])
+        tm.assert_series_equal(result, expected, check_index=False, check_names=False)
 
     def test_lookup_mixed(self, float_string_frame):
         df = float_string_frame
         rows = list(df.index) * len(df.columns)
         cols = list(df.columns) * len(df.index)
         result = df.lookup(rows, cols)
 
-        expected = np.array(
-            [df.loc[r, c] for r, c in zip(rows, cols)], dtype=np.object_
-        )
-        tm.assert_almost_equal(result, expected)
+        expected = Series([df.loc[r, c] for r, c in zip(rows, cols)], dtype=np.object_)
+        tm.assert_series_equal(result, expected, check_index=False, check_names=False)
 
     def test_lookup_bool(self):
         df = DataFrame(
@@ -1443,14 +1441,16 @@ def test_lookup_bool(self):
                 "mask_c": [False, True, False, True],
             }
         )
-        df["mask"] = df.lookup(df.index, "mask_" + df["label"])
+        df_mask = df.lookup(df.index, "mask_" + df["label"])
 
-        exp_mask = np.array(
+        exp_mask = Series(
             [df.loc[r, c] for r, c in zip(df.index, "mask_" + df["label"])]
         )
 
-        tm.assert_series_equal(df["mask"], Series(exp_mask, name="mask"))
-        assert df["mask"].dtype == np.bool_
+        tm.assert_series_equal(
+            df_mask, Series(exp_mask, name="mask"), check_index=False, check_names=False
+        )
+        assert df_mask.dtype == np.bool_
 
     def test_lookup_raises(self, float_frame):
         with pytest.raises(KeyError, match="'One or more row labels was not found'"):