Compile take functions for uint dtypes (pandas-dev#57632)

Author: phofl

pandas/_libs/algos.pyi

@@ -165,6 +165,15 @@ def take_1d_int32_float64(
 def take_1d_int64_int64(
     values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=...
 ) -> None: ...
+def take_1d_uint16_uint16(
+    values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=...
+) -> None: ...
+def take_1d_uint32_uint32(
+    values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=...
+) -> None: ...
+def take_1d_uint64_uint64(
+    values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=...
+) -> None: ...
 def take_1d_int64_float64(
     values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=...
 ) -> None: ...
@@ -225,6 +234,15 @@ def take_2d_axis0_int64_int64(
 def take_2d_axis0_int64_float64(
     values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=...
 ) -> None: ...
+def take_2d_axis0_uint16_uint16(
+    values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=...
+) -> None: ...
+def take_2d_axis0_uint32_uint32(
+    values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=...
+) -> None: ...
+def take_2d_axis0_uint64_uint64(
+    values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=...
+) -> None: ...
 def take_2d_axis0_float32_float32(
     values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=...
 ) -> None: ...
@@ -279,6 +297,15 @@ def take_2d_axis1_int32_float64(
 def take_2d_axis1_int64_int64(
     values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=...
 ) -> None: ...
+def take_2d_axis1_uint16_uint16(
+    values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=...
+) -> None: ...
+def take_2d_axis1_uint32_uint32(
+    values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=...
+) -> None: ...
+def take_2d_axis1_uint64_uint64(
+    values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=...
+) -> None: ...
 def take_2d_axis1_int64_float64(
     values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=...
 ) -> None: ...

pandas/_libs/algos_take_helper.pxi.in

@@ -15,6 +15,9 @@ WARNING: DO NOT edit .pxi FILE directly, .pxi is generated from .pxi.in
 dtypes = [
     ('uint8_t', 'uint8_t'),
     ('uint8_t', 'object'),
+    ('uint16_t', 'uint16_t'),
+    ('uint32_t', 'uint32_t'),
+    ('uint64_t', 'uint64_t'),
     ('int8_t', 'int8_t'),
     ('int8_t', 'int32_t'),
     ('int8_t', 'int64_t'),

pandas/core/array_algos/take.py

@@ -337,6 +337,10 @@ def wrapper(
     ("int32", "int64"): libalgos.take_1d_int32_int64,
     ("int32", "float64"): libalgos.take_1d_int32_float64,
     ("int64", "int64"): libalgos.take_1d_int64_int64,
+    ("uint8", "uint8"): libalgos.take_1d_bool_bool,
+    ("uint16", "int64"): libalgos.take_1d_uint16_uint16,
+    ("uint32", "int64"): libalgos.take_1d_uint32_uint32,
+    ("uint64", "int64"): libalgos.take_1d_uint64_uint64,
     ("int64", "float64"): libalgos.take_1d_int64_float64,
     ("float32", "float32"): libalgos.take_1d_float32_float32,
     ("float32", "float64"): libalgos.take_1d_float32_float64,
@@ -366,6 +370,10 @@ def wrapper(
     ("int32", "float64"): libalgos.take_2d_axis0_int32_float64,
     ("int64", "int64"): libalgos.take_2d_axis0_int64_int64,
     ("int64", "float64"): libalgos.take_2d_axis0_int64_float64,
+    ("uint8", "uint8"): libalgos.take_2d_axis0_bool_bool,
+    ("uint16", "uint16"): libalgos.take_2d_axis0_uint16_uint16,
+    ("uint32", "uint32"): libalgos.take_2d_axis0_uint32_uint32,
+    ("uint64", "uint64"): libalgos.take_2d_axis0_uint64_uint64,
     ("float32", "float32"): libalgos.take_2d_axis0_float32_float32,
     ("float32", "float64"): libalgos.take_2d_axis0_float32_float64,
     ("float64", "float64"): libalgos.take_2d_axis0_float64_float64,
@@ -398,6 +406,10 @@ def wrapper(
     ("int32", "float64"): libalgos.take_2d_axis1_int32_float64,
     ("int64", "int64"): libalgos.take_2d_axis1_int64_int64,
     ("int64", "float64"): libalgos.take_2d_axis1_int64_float64,
+    ("uint8", "uint8"): libalgos.take_2d_axis1_bool_bool,
+    ("uint16", "uint16"): libalgos.take_2d_axis1_uint16_uint16,
+    ("uint32", "uint32"): libalgos.take_2d_axis1_uint32_uint32,
+    ("uint64", "uint64"): libalgos.take_2d_axis1_uint64_uint64,
     ("float32", "float32"): libalgos.take_2d_axis1_float32_float32,
     ("float32", "float64"): libalgos.take_2d_axis1_float32_float64,
     ("float64", "float64"): libalgos.take_2d_axis1_float64_float64,