Skip to content

Commit 70fa75a

Browse files
jschendeljreback
authored andcommitted
TST: Fix dtype mismatch on 32bit in IntervalTree get_indexer test (#23468)
1 parent 498bdd4 commit 70fa75a

File tree

2 files changed

+46
-24
lines changed

2 files changed

+46
-24
lines changed

Diff for: pandas/_libs/intervaltree.pxi.in

+4-3
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,7 @@ cdef class IntervalTree(IntervalMixin):
105105
self.root.query(result, key)
106106
if not result.data.n:
107107
raise KeyError(key)
108-
return result.to_array()
108+
return result.to_array().astype('intp')
109109

110110
def _get_partial_overlap(self, key_left, key_right, side):
111111
"""Return all positions corresponding to intervals with the given side
@@ -155,7 +155,7 @@ cdef class IntervalTree(IntervalMixin):
155155
raise KeyError(
156156
'indexer does not intersect a unique set of intervals')
157157
old_len = result.data.n
158-
return result.to_array()
158+
return result.to_array().astype('intp')
159159

160160
def get_indexer_non_unique(self, scalar_t[:] target):
161161
"""Return the positions corresponding to intervals that overlap with
@@ -175,7 +175,8 @@ cdef class IntervalTree(IntervalMixin):
175175
result.append(-1)
176176
missing.append(i)
177177
old_len = result.data.n
178-
return result.to_array(), missing.to_array()
178+
return (result.to_array().astype('intp'),
179+
missing.to_array().astype('intp'))
179180

180181
def __repr__(self):
181182
return ('<IntervalTree[{dtype},{closed}]: '

Diff for: pandas/tests/indexes/interval/test_interval_tree.py

+42-21
Original file line numberDiff line numberDiff line change
@@ -49,44 +49,64 @@ def tree(request, leaf_size):
4949
class TestIntervalTree(object):
5050

5151
def test_get_loc(self, tree):
52-
tm.assert_numpy_array_equal(tree.get_loc(1),
53-
np.array([0], dtype='int64'))
54-
tm.assert_numpy_array_equal(np.sort(tree.get_loc(2)),
55-
np.array([0, 1], dtype='int64'))
52+
result = tree.get_loc(1)
53+
expected = np.array([0], dtype='intp')
54+
tm.assert_numpy_array_equal(result, expected)
55+
56+
result = np.sort(tree.get_loc(2))
57+
expected = np.array([0, 1], dtype='intp')
58+
tm.assert_numpy_array_equal(result, expected)
59+
5660
with pytest.raises(KeyError):
5761
tree.get_loc(-1)
5862

5963
def test_get_indexer(self, tree):
60-
tm.assert_numpy_array_equal(
61-
tree.get_indexer(np.array([1.0, 5.5, 6.5])),
62-
np.array([0, 4, -1], dtype='int64'))
64+
result = tree.get_indexer(np.array([1.0, 5.5, 6.5]))
65+
expected = np.array([0, 4, -1], dtype='intp')
66+
tm.assert_numpy_array_equal(result, expected)
67+
6368
with pytest.raises(KeyError):
6469
tree.get_indexer(np.array([3.0]))
6570

6671
def test_get_indexer_non_unique(self, tree):
6772
indexer, missing = tree.get_indexer_non_unique(
6873
np.array([1.0, 2.0, 6.5]))
69-
tm.assert_numpy_array_equal(indexer[:1],
70-
np.array([0], dtype='int64'))
71-
tm.assert_numpy_array_equal(np.sort(indexer[1:3]),
72-
np.array([0, 1], dtype='int64'))
73-
tm.assert_numpy_array_equal(np.sort(indexer[3:]),
74-
np.array([-1], dtype='int64'))
75-
tm.assert_numpy_array_equal(missing, np.array([2], dtype='int64'))
74+
75+
result = indexer[:1]
76+
expected = np.array([0], dtype='intp')
77+
tm.assert_numpy_array_equal(result, expected)
78+
79+
result = np.sort(indexer[1:3])
80+
expected = np.array([0, 1], dtype='intp')
81+
tm.assert_numpy_array_equal(result, expected)
82+
83+
result = np.sort(indexer[3:])
84+
expected = np.array([-1], dtype='intp')
85+
tm.assert_numpy_array_equal(result, expected)
86+
87+
result = missing
88+
expected = np.array([2], dtype='intp')
89+
tm.assert_numpy_array_equal(result, expected)
7690

7791
def test_duplicates(self, dtype):
7892
left = np.array([0, 0, 0], dtype=dtype)
7993
tree = IntervalTree(left, left + 1)
80-
tm.assert_numpy_array_equal(np.sort(tree.get_loc(0.5)),
81-
np.array([0, 1, 2], dtype='int64'))
94+
95+
result = np.sort(tree.get_loc(0.5))
96+
expected = np.array([0, 1, 2], dtype='intp')
97+
tm.assert_numpy_array_equal(result, expected)
8298

8399
with pytest.raises(KeyError):
84100
tree.get_indexer(np.array([0.5]))
85101

86102
indexer, missing = tree.get_indexer_non_unique(np.array([0.5]))
87-
tm.assert_numpy_array_equal(np.sort(indexer),
88-
np.array([0, 1, 2], dtype='int64'))
89-
tm.assert_numpy_array_equal(missing, np.array([], dtype='int64'))
103+
result = np.sort(indexer)
104+
expected = np.array([0, 1, 2], dtype='intp')
105+
tm.assert_numpy_array_equal(result, expected)
106+
107+
result = missing
108+
expected = np.array([], dtype='intp')
109+
tm.assert_numpy_array_equal(result, expected)
90110

91111
def test_get_loc_closed(self, closed):
92112
tree = IntervalTree([0], [1], closed=closed)
@@ -96,8 +116,9 @@ def test_get_loc_closed(self, closed):
96116
with pytest.raises(KeyError):
97117
tree.get_loc(p)
98118
else:
99-
tm.assert_numpy_array_equal(tree.get_loc(p),
100-
np.array([0], dtype='int64'))
119+
result = tree.get_loc(p)
120+
expected = np.array([0], dtype='intp')
121+
tm.assert_numpy_array_equal(result, expected)
101122

102123
@pytest.mark.parametrize('leaf_size', [
103124
skipif_32bit(1), skipif_32bit(10), skipif_32bit(100), 10000])

0 commit comments

Comments
 (0)