Repository: pandas-dev/pandas
Pull request: BUG: Perform i8 conversion for datetimelike IntervalTree queries
- [X] closes #20636
- [X] tests added / passed
- [X] passes `git diff upstream/master -u -- "*.py" | flake8 --diff`
- [X] whatsnew entry

I've seen a few complaints about this on StackOverflow, so wanted to get this into 0.24.0 before I try implementing larger changes like the new behavior specs and interval accessor.

Problem statement:
BUG: indexing in datetime IntervalIndex with duplicate values fails
#### Code Sample
```python
import pandas as pd
start_end_idx = pd.IntervalIndex.from_arrays(start, end)
start_end_idx.name = 'port_time'
mm = mm.set_index(['imo','port_id', start_end_idx])
global_bins = mm.index
global_bins.get_loc((8801527, 5610.0, pd.to_datetime('2016-10-16 16:57:44.316')))
```
```
/anaconda3/envs/mariquant/lib/python3.6/site-packages/pandas/core/indexes/interval.py in _engine(self)
    283     @cache_readonly
    284     def _engine(self):
--> 285         return IntervalTree(self.left, self.right, closed=self.closed)
    286 
    287     @property

pandas/_libs/intervaltree.pxi in pandas._libs.interval.IntervalTree.__init__()

KeyError: ('datetime64[ns]', 'right')
```
#### Problem description
It is impossible to use non unique IntervalIndex with datetimes as start-end points. Both index.get_loc and DataFrame.loc produce the same error in IntervalTree.

#### Expected Output

#### Output of ``pd.show_versions()``

<details>

[paste the output of ``pd.show_versions()`` here below this line]

INSTALLED VERSIONS
------------------
commit: None

pandas: 0.22.0
pytest: None
pip: 9.0.1
setuptools: 38.5.1
Cython: None
numpy: 1.14.2
scipy: 1.0.0
pyarrow: None
xarray: 0.10.2
IPython: 6.2.1
sphinx: 1.7.1
patsy: 0.5.0
dateutil: 2.7.0
pytz: 2018.3
blosc: None
bottleneck: 1.2.1
tables: None
numexpr: None
feather: None
matplotlib: 2.2.2
openpyxl: None
xlrd: None
xlwt: None
xlsxwriter: None
lxml: None
bs4: None
html5lib: 1.0.1
sqlalchemy: 1.2.5
pymysql: None
psycopg2: 2.7.4 (dt dec pq3 ext lo64)
jinja2: 2.10
s3fs: None
fastparquet: None
pandas_gbq: None
pandas_datareader: None

</details>

Hints from issue discussion:
@avnovikov Can you provide a reproducible example? (in this case, provide example code for `start`, `end` and `mm`) See also http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
I can reproduce the error on master with a duplicated `IntervalIndex`:
```python-traceback
In [2]: ii = pd.interval_range(pd.Timestamp('20180101'), periods=2).repeat(2)

In [3]: ii
Out[3]:
IntervalIndex([(2018-01-01, 2018-01-02], (2018-01-01, 2018-01-02], (2018-01-02, 2018-01-03], (2018-01-02, 2018-01-03]]
              closed='right',
              dtype='interval[datetime64[ns]]')

In [4]: ii.get_loc(pd.Timestamp('20180102'))
---------------------------------------------------------------------------
KeyError: ('datetime64[ns]', 'right')
```
OK, and for numeric intervals, this works correctly:

```
In [30]: ii = pd.interval_range(1, periods=2).repeat(2)

In [31]: ii
Out[31]: 
IntervalIndex([(1, 2], (1, 2], (2, 3], (2, 3]]
              closed='right',
              dtype='interval[int64]')

In [33]:  ii.get_loc(2)
Out[33]: array([0, 1])
```
I suspect this will fail for dtypes other than `float32`/`float64`/`int32`/`int64` when the `IntervalIndex` is overlapping or not monotonic (the non-overlapping monotonic case doesn't require using `IntervalTree`). It looks like `IntervalTree` was only implemented for the four dtypes I mentioned earlier:

https://github.com/pandas-dev/pandas/blob/4e6aa1c012d91d663edcfd13658a35f88a49da30/pandas/_libs/intervaltree.pxi.in#L202-L208

Using `timedelta64[ns]` fails:
```python-traceback
In [2]: ii = pd.interval_range(pd.Timedelta('0 days'), periods=2).repeat(2)

In [3]: ii
Out[3]:
IntervalIndex([(0 days 00:00:00, 1 days 00:00:00], (0 days 00:00:00, 1 days 00:00:00], (1 days 00:00:00, 2 days 00:00:00], (1 days 00:00:00, 2 days 00:00:00]]
              closed='right',
              dtype='interval[timedelta64[ns]]')

In [4]: ii.get_loc(pd.Timedelta('1 day'))
---------------------------------------------------------------------------
KeyError: ('timedelta64[ns]', 'right')
```

Using `uint64` fails:
```python-traceback
In [11]: ii = pd.interval_range(1, periods=2).repeat(2).astype('interval[uint64]')

In [12]: ii
Out[12]:
IntervalIndex([(1, 2], (1, 2], (2, 3], (2, 3]]
              closed='right',
              dtype='interval[uint64]')

In [13]: ii.get_loc(1.5)
---------------------------------------------------------------------------
KeyError: ('uint64', 'right')
```

I'm guessing for datetimelike we'll need to do i8 conversion?  Or is there a way to add that directly?  I think `uint64` can be directly added as a dtype.  Any other dtypes I'm forgetting?

Commit message: BUG: Perform i8 conversion for datetimelike IntervalTree queries

Patch to review:
diff --git a/doc/source/whatsnew/v0.24.0.txt b/doc/source/whatsnew/v0.24.0.txt
index f246ebad3aa2c..42c4134437ff6 100644
--- a/doc/source/whatsnew/v0.24.0.txt
+++ b/doc/source/whatsnew/v0.24.0.txt
@@ -755,7 +755,7 @@ Interval
 - Bug in the :class:`IntervalIndex` constructor where the ``closed`` parameter did not always override the inferred ``closed`` (:issue:`19370`)
 - Bug in the ``IntervalIndex`` repr where a trailing comma was missing after the list of intervals (:issue:`20611`)
 - Bug in :class:`Interval` where scalar arithmetic operations did not retain the ``closed`` value (:issue:`22313`)
--
+- Bug in :class:`IntervalIndex` where indexing with datetime-like values raised a ``KeyError`` (:issue:`20636`)
 
 Indexing
 ^^^^^^^^
diff --git a/pandas/core/indexes/interval.py b/pandas/core/indexes/interval.py
index f72f87aeb2af6..9b8570796a455 100644
--- a/pandas/core/indexes/interval.py
+++ b/pandas/core/indexes/interval.py
@@ -6,12 +6,14 @@
 
 from pandas.compat import add_metaclass
 from pandas.core.dtypes.missing import isna
-from pandas.core.dtypes.cast import find_common_type, maybe_downcast_to_dtype
+from pandas.core.dtypes.cast import (
+    find_common_type, maybe_downcast_to_dtype, infer_dtype_from_scalar)
 from pandas.core.dtypes.common import (
     ensure_platform_int,
     is_list_like,
     is_datetime_or_timedelta_dtype,
     is_datetime64tz_dtype,
+    is_dtype_equal,
     is_integer_dtype,
     is_float_dtype,
     is_interval_dtype,
@@ -19,7 +21,8 @@
     is_scalar,
     is_float,
     is_number,
-    is_integer)
+    is_integer,
+    needs_i8_conversion)
 from pandas.core.indexes.base import (
     Index, ensure_index,
     default_pprint, _index_shared_docs)
@@ -29,8 +32,8 @@
     Interval, IntervalMixin, IntervalTree,
 )
 
-from pandas.core.indexes.datetimes import date_range
-from pandas.core.indexes.timedeltas import timedelta_range
+from pandas.core.indexes.datetimes import date_range, DatetimeIndex
+from pandas.core.indexes.timedeltas import timedelta_range, TimedeltaIndex
 from pandas.core.indexes.multi import MultiIndex
 import pandas.core.common as com
 from pandas.util._decorators import cache_readonly, Appender
@@ -192,7 +195,9 @@ def _isnan(self):
 
     @cache_readonly
     def _engine(self):
-        return IntervalTree(self.left, self.right, closed=self.closed)
+        left = self._maybe_convert_i8(self.left)
+        right = self._maybe_convert_i8(self.right)
+        return IntervalTree(left, right, closed=self.closed)
 
     def __contains__(self, key):
         """
@@ -514,6 +519,41 @@ def _maybe_cast_indexed(self, key):
 
         return key
 
+    def _maybe_convert_i8(self, key):
+        if isinstance(key, Interval):
+            if not isinstance(key.left, (Timestamp, Timedelta)):
+                return key
+            left = self._maybe_convert_i8(key.left)
+            right = self._maybe_convert_i8(key.right)
+            return Interval(left, right, key.closed)
+        elif isinstance(key, (IntervalIndex, IntervalArray)):
+            if not needs_i8_conversion(key.left):
+                return key
+            left = self._maybe_convert_i8(key.left)
+            right = self._maybe_convert_i8(key.right)
+            return IntervalIndex.from_arrays(left, right, key.closed)
+        elif is_list_like(key) and not isinstance(key, Index):
+            result = self._maybe_convert_i8(ensure_index(key))
+            if result[0] == key[0]:
+                # return the list-like key if no conversion
+                return key
+            return result
+
+        subtype = self.dtype.subtype
+        msg = ('Cannot index an IntervalIndex of subtype {subtype} with '
+               'values of dtype {other}')
+        if isinstance(key, (Timestamp, Timedelta)):
+            key_dtype, key_i8 = infer_dtype_from_scalar(key, pandas_dtype=True)
+            if not is_dtype_equal(subtype, key_dtype):
+                raise ValueError(msg.format(subtype=subtype, other=key_dtype))
+            return key_i8
+        elif isinstance(key, (DatetimeIndex, TimedeltaIndex)):
+            if not is_dtype_equal(subtype, key.dtype):
+                raise ValueError(msg.format(subtype=subtype, other=key.dtype))
+            return Index(key.asi8)
+
+        return key
+
     def _check_method(self, method):
         if method is None:
             return
@@ -648,6 +688,7 @@ def get_loc(self, key, method=None):
 
         else:
             # use the interval tree
+            key = self._maybe_convert_i8(key)
             if isinstance(key, Interval):
                 left, right = _get_interval_closed_bounds(key)
                 return self._engine.get_loc_interval(left, right)
@@ -711,8 +752,10 @@ def _get_reindexer(self, target):
         """
 
         # find the left and right indexers
-        lindexer = self._engine.get_indexer(target.left.values)
-        rindexer = self._engine.get_indexer(target.right.values)
+        left = self._maybe_convert_i8(target.left)
+        right = self._maybe_convert_i8(target.right)
+        lindexer = self._engine.get_indexer(left.values)
+        rindexer = self._engine.get_indexer(right.values)
 
         # we want to return an indexer on the intervals
         # however, our keys could provide overlapping of multiple
diff --git a/pandas/tests/indexes/interval/test_interval.py b/pandas/tests/indexes/interval/test_interval.py
index 71f56c5bc1164..559b03d51407d 100644
--- a/pandas/tests/indexes/interval/test_interval.py
+++ b/pandas/tests/indexes/interval/test_interval.py
@@ -1,7 +1,9 @@
 from __future__ import division
 
+from itertools import permutations
 import pytest
 import numpy as np
+import re
 from pandas import (
     Interval, IntervalIndex, Index, isna, notna, interval_range, Timestamp,
     Timedelta, date_range, timedelta_range)
@@ -498,6 +500,48 @@ def test_get_loc_length_one(self, item, closed):
         result = index.get_loc(item)
         assert result == 0
 
+    # Make consistent with test_interval_new.py (see #16316, #16386)
+    @pytest.mark.parametrize('breaks', [
+        date_range('20180101', periods=4),
+        date_range('20180101', periods=4, tz='US/Eastern'),
+        timedelta_range('0 days', periods=4)], ids=lambda x: str(x.dtype))
+    def test_get_loc_datetimelike_nonoverlapping(self, breaks):
+        # GH 20636
+        # nonoverlapping = IntervalIndex method and no i8 conversion
+        index = IntervalIndex.from_breaks(breaks)
+
+        value = index[0].mid
+        result = index.get_loc(value)
+        expected = 0
+        assert result == expected
+
+        interval = Interval(index[0].left, index[1].right)
+        result = index.get_loc(interval)
+        expected = slice(0, 2)
+        assert result == expected
+
+    # Make consistent with test_interval_new.py (see #16316, #16386)
+    @pytest.mark.parametrize('arrays', [
+        (date_range('20180101', periods=4), date_range('20180103', periods=4)),
+        (date_range('20180101', periods=4, tz='US/Eastern'),
+         date_range('20180103', periods=4, tz='US/Eastern')),
+        (timedelta_range('0 days', periods=4),
+         timedelta_range('2 days', periods=4))], ids=lambda x: str(x[0].dtype))
+    def test_get_loc_datetimelike_overlapping(self, arrays):
+        # GH 20636
+        # overlapping = IntervalTree method with i8 conversion
+        index = IntervalIndex.from_arrays(*arrays)
+
+        value = index[0].mid + Timedelta('12 hours')
+        result = np.sort(index.get_loc(value))
+        expected = np.array([0, 1], dtype='int64')
+        assert tm.assert_numpy_array_equal(result, expected)
+
+        interval = Interval(index[0].left, index[1].right)
+        result = np.sort(index.get_loc(interval))
+        expected = np.array([0, 1, 2], dtype='int64')
+        assert tm.assert_numpy_array_equal(result, expected)
+
     # To be removed, replaced by test_interval_new.py (see #16316, #16386)
     def test_get_indexer(self):
         actual = self.index.get_indexer([-1, 0, 0.5, 1, 1.5, 2, 3])
@@ -555,6 +599,97 @@ def test_get_indexer_length_one(self, item, closed):
         expected = np.array([0] * len(item), dtype='intp')
         tm.assert_numpy_array_equal(result, expected)
 
+    # Make consistent with test_interval_new.py (see #16316, #16386)
+    @pytest.mark.parametrize('arrays', [
+        (date_range('20180101', periods=4), date_range('20180103', periods=4)),
+        (date_range('20180101', periods=4, tz='US/Eastern'),
+         date_range('20180103', periods=4, tz='US/Eastern')),
+        (timedelta_range('0 days', periods=4),
+         timedelta_range('2 days', periods=4))], ids=lambda x: str(x[0].dtype))
+    def test_get_reindexer_datetimelike(self, arrays):
+        # GH 20636
+        index = IntervalIndex.from_arrays(*arrays)
+        tuples = [(index[0].left, index[0].left + pd.Timedelta('12H')),
+                  (index[-1].right - pd.Timedelta('12H'), index[-1].right)]
+        target = IntervalIndex.from_tuples(tuples)
+
+        result = index._get_reindexer(target)
+        expected = np.array([0, 3], dtype='int64')
+        tm.assert_numpy_array_equal(result, expected)
+
+    @pytest.mark.parametrize('breaks', [
+        date_range('20180101', periods=4),
+        date_range('20180101', periods=4, tz='US/Eastern'),
+        timedelta_range('0 days', periods=4)], ids=lambda x: str(x.dtype))
+    def test_maybe_convert_i8(self, breaks):
+        # GH 20636
+        index = IntervalIndex.from_breaks(breaks)
+
+        # intervalindex
+        result = index._maybe_convert_i8(index)
+        expected = IntervalIndex.from_breaks(breaks.asi8)
+        tm.assert_index_equal(result, expected)
+
+        # interval
+        interval = Interval(breaks[0], breaks[1])
+        result = index._maybe_convert_i8(interval)
+        expected = Interval(breaks[0].value, breaks[1].value)
+        assert result == expected
+
+        # datetimelike index
+        result = index._maybe_convert_i8(breaks)
+        expected = Index(breaks.asi8)
+        tm.assert_index_equal(result, expected)
+
+        # datetimelike scalar
+        result = index._maybe_convert_i8(breaks[0])
+        expected = breaks[0].value
+        assert result == expected
+
+        # list-like of datetimelike scalars
+        result = index._maybe_convert_i8(list(breaks))
+        expected = Index(breaks.asi8)
+        tm.assert_index_equal(result, expected)
+
+    @pytest.mark.parametrize('breaks', [
+        np.arange(5, dtype='int64'),
+        np.arange(5, dtype='float64')], ids=lambda x: str(x.dtype))
+    def test_maybe_convert_i8_numeric(self, breaks):
+        # GH 20636
+        index = IntervalIndex.from_breaks(breaks)
+        numeric_keys = [
+            IntervalIndex.from_breaks(breaks),
+            Interval(breaks[0], breaks[1]),
+            breaks,
+            breaks[0],
+            list(breaks)]
+
+        # no conversion occurs for numeric
+        for key in numeric_keys:
+            result = index._maybe_convert_i8(key)
+            assert result is key
+
+    @pytest.mark.parametrize('breaks1, breaks2', permutations([
+        date_range('20180101', periods=4),
+        date_range('20180101', periods=4, tz='US/Eastern'),
+        timedelta_range('0 days', periods=4)], 2), ids=lambda x: str(x.dtype))
+    def test_maybe_convert_i8_errors(self, breaks1, breaks2):
+        # GH 20636
+        index = IntervalIndex.from_breaks(breaks1)
+        invalid_keys = [
+            IntervalIndex.from_breaks(breaks2),
+            Interval(breaks2[0], breaks2[1]),
+            breaks2,
+            breaks2[0],
+            list(breaks2)]
+
+        msg = ('Cannot index an IntervalIndex of subtype {dtype1} with '
+               'values of dtype {dtype2}')
+        msg = re.escape(msg.format(dtype1=breaks1.dtype, dtype2=breaks2.dtype))
+        for key in invalid_keys:
+            with tm.assert_raises_regex(ValueError, msg):
+                index._maybe_convert_i8(key)
+
     # To be removed, replaced by test_interval_new.py (see #16316, #16386)
     def test_contains(self):
         # Only endpoints are valid.


Code review comments:
diff_hunk: @@ -514,6 +519,41 @@ def _maybe_cast_indexed(self, key):
 
         return key
 
+    def _maybe_convert_i8(self, key):
+        if isinstance(key, Interval):
path: pandas/core/indexes/interval.py
text: @user1:
can you document this a bit, and can you simplify at all?

@author:
Done. Look better?