Fix indexing tests

Author: philippjfr

holoviews/core/data/pandas.py

@@ -2,6 +2,7 @@
 import pandas as pd
 from packaging.version import Version
 from pandas.api.types import is_numeric_dtype
+from pandas.core.dtypes.common import ensure_int64
 
 from .. import util
 from ..dimension import Dimension, dimension_name
@@ -351,11 +352,36 @@ def sort(cls, dataset, by=None, reverse=False):
             return dataset.data.sort(columns=cols, ascending=not reverse)
         return dataset.data.sort_values(by=cols, ascending=not reverse)
 
+    @classmethod
+    def sorted_index(cls, df):
+        if hasattr(df.index, 'is_lexsorted'):
+            return df.index.is_lexsorted()
+        return df.index.is_monotonic_increasing
+
+    @classmethod
+    def sort_depth(cls, df):
+        try:
+            from pandas._libs.algos import is_lexsorted
+        except Exception:
+            return 0
+        int64_codes = [ensure_int64(level_codes) for level_codes in df.index.codes]
+        for k in range(df.index.nlevels, 0, -1):
+            if is_lexsorted(int64_codes[:k]):
+                return k
+        return 0
+
     @classmethod
     def index_selection(cls, df, selection):
+        indexes = cls.indexes(df)
+        nindex = len(indexes)
+        sorted_index = cls.sorted_index(df)
+        if sorted_index:
+            depth = df.index.nlevels
+        else:
+            depth = cls.sort_depth(df)
         index_sel = {}
         skip_index = True
-        for idx in cls.indexes(df):
+        for level, idx in enumerate(indexes):
             if idx not in selection:
                 index_sel[idx] = slice(None, None)
                 continue
@@ -365,18 +391,26 @@ def index_selection(cls, df, selection):
                 sel = slice(*sel)
             elif not isinstance(sel, (list, slice)):
                 sel = [sel]
+            if isinstance(sel, slice) and nindex > 1 and not sorted_index and level>depth:
+                # If the index is not monotonic we cannot slice
+                return {}
             index_sel[idx] = sel
+        print(index_sel)
         return {} if skip_index else index_sel
 
     @classmethod
     def select(cls, dataset, selection_mask=None, **selection):
         df = dataset.data
         if selection_mask is None:
             if index_sel:= cls.index_selection(df, selection):
-                if len(index_sel) == 1:
-                    df = df[next(iter(index_sel.values()))]
-                else:
-                    df = df.loc[tuple(index_sel.values()), :]
+                try:
+                    if len(index_sel) == 1:
+                        df = df[next(iter(index_sel.values()))]
+                    else:
+                        df = df.loc[tuple(index_sel.values()), :]
+                except KeyError:
+                    # If index lookup fails we fall back to boolean indexing
+                    index_sel = {}
             column_sel = {k: v for k, v in selection.items() if k not in index_sel}
             if column_sel:
                 selection_mask = cls.select_mask(dataset, column_sel)

holoviews/tests/core/data/test_pandasinterface.py

@@ -222,6 +222,13 @@ def test_index_aggregate(self):
         agg = ds.aggregate("number", function=np.mean, spreadfn=np.var)
         pd.testing.assert_frame_equal(agg.data, expected)
 
+    def test_index_select_monotonic(self):
+        ds = Dataset(self.df, kdims=["number", "color"])
+        selected = ds.select(number=1)
+        expected = pd.DataFrame({'color': ['red', 'blue'], 'values': [0, 1], 'number': [1, 1]}).set_index(['number', 'color'])
+        assert isinstance(selected.data.index, pd.MultiIndex)
+        pd.testing.assert_frame_equal(selected.data, expected)
+
     def test_index_select(self):
         ds = Dataset(self.df, kdims=["number", "color"])
         selected = ds.select(number=1)
@@ -307,31 +314,29 @@ def test_sort(self):
         np.testing.assert_array_equal(sorted_ds.dimension_values("values"), [1, 3, 0, 2])
         np.testing.assert_array_equal(sorted_ds.dimension_values("number"), [1, 2, 1, 2])
 
-    def test_select(self):
-        ds = Dataset(self.df, kdims=["number", "color"])
+    def test_select_monotonic(self):
+        ds = Dataset(self.df.sort_index(), kdims=["number", "color"])
         selected = ds.select(color="red")
         pd.testing.assert_frame_equal(selected.data, self.df.iloc[[0, 2], :])
 
         selected = ds.select(number=1, color='red')
         assert selected == 0
 
-    @pytest.mark.xfail(reason="Not working")
     def test_select_not_monotonic(self):
         frame = pd.DataFrame({"number": [1, 1, 2, 2], "color": [2, 1, 2, 1]})
         index = pd.MultiIndex.from_frame(frame, names=frame.columns)
         df = pd.DataFrame(range(4), index=index, columns=["values"])
         ds = Dataset(df, kdims=list(frame.columns))
 
         data = ds.select(color=slice(2, 3)).data
-        expected = pd.DataFrame({"number": [1, 2], "color": [2, 2], "values": [1, 3]}).set_index(['number', 'color'])
+        expected = pd.DataFrame({"number": [1, 2], "color": [2, 2], "values": [0, 2]}).set_index(['number', 'color'])
         pd.testing.assert_frame_equal(data, expected)
 
-    @pytest.mark.xfail(reason="Not working")
     def test_select_not_in_index(self):
         ds = Dataset(self.df, kdims=["number", "color"])
         selected = ds.select(number=[2, 3])
-        expected = ds.select(number=2)
-        pd.testing.assert_frame_equal(selected.data, expected.data)
+        expected = self.df.loc[[2]]
+        pd.testing.assert_frame_equal(selected.data, expected)
 
     def test_sample(self):
         ds = Dataset(self.df, kdims=["number", "color"])