Add last tests and updates for iloc

holoviz · Feb 21, 2024 · 5ec315a · 5ec315a
1 parent 03af717
commit 5ec315a
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 6 deletions.
diff --git a/holoviews/core/data/pandas.py b/holoviews/core/data/pandas.py
@@ -486,12 +486,12 @@ def iloc(cls, dataset, index):
         columns = list(data.columns)
         id_cols = [columns.index(c) for c in cols if c not in indexes]
         if not id_cols:
-            if scalar:
-                return data.index.values[rows[0]]
-            elif len(indexes) > 1:
-                return data.index.to_frame()[cols].iloc[rows].reset_index(drop=True)
+            if len(indexes) > 1:
+                data = data.index.to_frame()[cols].iloc[rows].reset_index(drop=True)
+                data = data.values.ravel()[0] if scalar else data
             else:
-                return data.index[rows]
+                data = data.index.values[rows[0]] if scalar else data.index[rows]
+            return data
         if scalar:
             return data.iloc[rows[0], id_cols[0]]
         return data.iloc[rows, id_cols]

diff --git a/holoviews/tests/core/data/test_pandasinterface.py b/holoviews/tests/core/data/test_pandasinterface.py
@@ -220,7 +220,13 @@ def test_index_select(self):
         assert isinstance(selected.data.index, pd.MultiIndex)
         pd.testing.assert_frame_equal(selected.data, expected)
 
-    def test_index_iloc_slice_scalar(self):
+    def test_index_iloc_scalar_scalar_only_index(self):
+        ds = Dataset(self.df, kdims=["number", "color"])
+        selected = ds.iloc[0, 0]
+        expected = 1
+        assert selected == expected
+
+    def test_index_iloc_slice_scalar_only_index(self):
         ds = Dataset(self.df, kdims=["number", "color"])
         selected = ds.iloc[:, 0]
         expected = self.df.reset_index()[["number"]]
@@ -237,3 +243,27 @@ def test_index_iloc_scalar_slice_only_index(self):
         selected = ds.iloc[0, :2]
         expected = pd.DataFrame({"number": 1, "color": "red"}, index=[0])
         pd.testing.assert_frame_equal(selected.data, expected)
+
+    def test_index_iloc_scalar_scalar(self):
+        ds = Dataset(self.df, kdims=["number", "color"])
+        selected = ds.iloc[0, 2]
+        expected = 0
+        assert selected == expected
+
+    def test_index_iloc_slice_scalar(self):
+        ds = Dataset(self.df, kdims=["number", "color"])
+        selected = ds.iloc[:, 2]
+        expected = self.df.iloc[:, [0]]
+        pd.testing.assert_frame_equal(selected.data, expected)
+
+    def test_index_iloc_slice_slice(self):
+        ds = Dataset(self.df, kdims=["number", "color"])
+        selected = ds.iloc[:, :3]
+        expected = self.df.iloc[:, [0]]
+        pd.testing.assert_frame_equal(selected.data, expected)
+
+    def test_index_iloc_scalar_slice(self):
+        ds = Dataset(self.df, kdims=["number", "color"])
+        selected = ds.iloc[0, :3]
+        expected = self.df.iloc[[0], [0]]
+        pd.testing.assert_frame_equal(selected.data, expected)