From e68a9eec173c6ea3516efa6a5b71da1ff1665cb5 Mon Sep 17 00:00:00 2001
From: Amanda Potts <ajpotts@users.noreply.github.com>
Date: Tue, 21 May 2024 14:54:12 -0400
Subject: [PATCH] Closes #3217: MultiIndex.get_level_values

---
 PROTO_tests/tests/index_test.py | 38 +++++++++++++++++++++++++++++++++
 arkouda/index.py                | 22 +++++++++++++++++++
 tests/index_test.py             | 38 +++++++++++++++++++++++++++++++++
 3 files changed, 98 insertions(+)

diff --git a/PROTO_tests/tests/index_test.py b/PROTO_tests/tests/index_test.py
index 1db6cdb5f6..c4f032b889 100644
--- a/PROTO_tests/tests/index_test.py
+++ b/PROTO_tests/tests/index_test.py
@@ -3,6 +3,7 @@
 
 import arkouda as ak
 from arkouda.dtypes import dtype
+from arkouda.index import Index
 from arkouda.pdarrayclass import pdarray
 
 
@@ -123,6 +124,43 @@ def test_eq(self):
         with pytest.raises(TypeError):
             i.equals("string")
 
+    def test_get_level_values(self):
+        m = ak.MultiIndex(
+            [ak.arange(3), ak.arange(3) * -1, ak.array(["a", 'b","c', "d"])],
+            names=["col1", "col2", "col3"],
+        )
+
+        i1 = Index(ak.arange(3), name="col1")
+        self.assert_equal(m.get_level_values(0), i1)
+        self.assert_equal(m.get_level_values("col1"), i1)
+
+        i2 = Index(ak.arange(3) * -1, name="col2")
+        self.assert_equal(m.get_level_values(1), i2)
+        self.assert_equal(m.get_level_values("col2"), i2)
+
+        i3 = Index(ak.array(["a", 'b","c', "d"]), name="col3")
+        self.assert_equal(m.get_level_values(2), i3)
+        self.assert_equal(m.get_level_values("col3"), i3)
+
+        with pytest.raises(ValueError):
+            m.get_level_values("col4")
+
+        #   Test when names=None
+        m2 = ak.MultiIndex(
+            [ak.arange(3), ak.arange(3) * -1, ak.array(["a", 'b","c', "d"])],
+        )
+        i4 = Index(ak.arange(3))
+        self.assert_equal(m2.get_level_values(0), i4)
+
+        with pytest.raises(RuntimeError):
+            m2.get_level_values("col")
+
+        with pytest.raises(ValueError):
+            m2.get_level_values(m2.nlevels)
+
+        with pytest.raises(ValueError):
+            m2.get_level_values(-1 * m2.nlevels)
+
     @pytest.mark.parametrize("size", pytest.prob_size)
     def test_memory_usage(self, size):
         from arkouda.dtypes import BigInt
diff --git a/arkouda/index.py b/arkouda/index.py
index 16719d7785..b5920d100b 100644
--- a/arkouda/index.py
+++ b/arkouda/index.py
@@ -1030,6 +1030,28 @@ def nlevels(self) -> int:
     def inferred_type(self) -> str:
         return "mixed"
 
+    def get_level_values(self, level: Union[str, int]):
+        if isinstance(level, str):
+            if self.names is None:
+                raise RuntimeError("Cannot get level values because Index.names is None.")
+            elif level not in self.names:
+                raise ValueError(
+                    f'Cannot get level values because level "{level}" is not in Index.names.'
+                )
+            elif isinstance(self.names, list) and level in self.names:
+                level = self.names.index(level)
+
+        if isinstance(level, int) and abs(level) < self.nlevels:
+            name = None
+            if isinstance(self.names, list) and level in self.names:
+                name = self.names[level]
+            return Index(self.levels[level], name=name)
+        else:
+            raise ValueError(
+                "Cannot get level values because level must be a string in names or "
+                "an integer with absolute value less than the number of levels."
+            )
+
     def memory_usage(self, unit="B"):
         """
         Return the memory usage of the MultiIndex levels.
diff --git a/tests/index_test.py b/tests/index_test.py
index febfa202b3..48e0b3fc11 100644
--- a/tests/index_test.py
+++ b/tests/index_test.py
@@ -8,6 +8,7 @@
 
 from arkouda import io_util
 from arkouda.dtypes import dtype
+from arkouda.index import Index
 from arkouda.pdarrayclass import pdarray
 
 
@@ -123,6 +124,43 @@ def test_eq(self):
         i5 = ak.Index(["a", "b", "c", "d"], allow_list=True)
         assert not i4.equals(i5)
 
+    def test_get_level_values(self):
+        m = ak.MultiIndex(
+            [ak.arange(3), ak.arange(3) * -1, ak.array(["a", 'b","c', "d"])],
+            names=["col1", "col2", "col3"],
+        )
+
+        i1 = Index(ak.arange(3), name="col1")
+        self.assert_equal(m.get_level_values(0), i1)
+        self.assert_equal(m.get_level_values("col1"), i1)
+
+        i2 = Index(ak.arange(3) * -1, name="col2")
+        self.assert_equal(m.get_level_values(1), i2)
+        self.assert_equal(m.get_level_values("col2"), i2)
+
+        i3 = Index(ak.array(["a", 'b","c', "d"]), name="col3")
+        self.assert_equal(m.get_level_values(2), i3)
+        self.assert_equal(m.get_level_values("col3"), i3)
+
+        with self.assertRaises(ValueError):
+            m.get_level_values("col4")
+
+        #   Test when names=None
+        m2 = ak.MultiIndex(
+            [ak.arange(3), ak.arange(3) * -1, ak.array(["a", 'b","c', "d"])],
+        )
+        i4 = Index(ak.arange(3))
+        self.assert_equal(m2.get_level_values(0), i4)
+
+        with self.assertRaises(RuntimeError):
+            m2.get_level_values("col")
+
+        with self.assertRaises(ValueError):
+            m2.get_level_values(m2.nlevels)
+
+        with self.assertRaises(ValueError):
+            m2.get_level_values(-1 * m2.nlevels)
+
     def test_memory_usage(self):
         from arkouda.dtypes import BigInt
         from arkouda.index import Index, MultiIndex