PERF: casting loc to labels dtype before searchsorted

[jorisvandenbossche]

Intrigued by the profiling results of the below example (Multi-index loc indexing, based on the example in #14549), where searchsorted seemed to take the majority of the computation time.
And it seems that searchsorted casts both inputs (in this case labels and loc) to a common dtype, and the labels of the MultiIndex were in this case int16, while loc (output from Index.get_loc) is a python int.

By casting loc to the dtype of labels, the specific example gets a ca 20 x speed improvement

df = pd.DataFrame({'a': np.random.randn(500*5000)}, index=pd.MultiIndex.from_product([date_range("2014-01-01", periods=500), range(5000)]))
dt = pd.Timestamp('2015-01-01')
%timeit df.loc[dt]

On master:

In [3]: %timeit df.loc[dt]
The slowest run took 5.70 times longer than the fastest. This could mean that an intermediate result is being cached.
100 loops, best of 3: 9.39 ms per loop

with this PR:

In [3]: %timeit df.loc[dt]
The slowest run took 122.51 times longer than the fastest. This could mean that an intermediate result is being cached.
1000 loops, best of 3: 422 µs per loop

Just putting it here as a showcase.
Actual PR probably needs some more work (other places where this can be done, can loc ever be out of bound for that dtype?, benchmarks, ..)

[jreback]

maybe should do this like

loc, orig_loc = lables.dtype.type(loc), loc
if  loc != orig_loc:
    loc = orig_loc

[jorisvandenbossche]

Yeah, I am not sure how much checking would be needed here.

My understanding is that it probably is not needed in this case, as loc is coming from loc = level_index.get_loc(key), and labels and level_index are from the same MultiIndex. So I would assume that get_loc can only return an existing label, and so should fit in the dtype of labels?

(but probably also not that a perf issue to do the check)

[jreback]

that sounds right; if the test suite passes, prob ok!

[jorisvandenbossche]

Apparently not :-)
But it's only a partial indexing test that fails. So loc could also be a slice, and it is expected in partial indexing that this raises an error in searchsorted, but now it raises already a slightly different error in dtype.type(loc). So that can easily be solved with:

                try:
                    loc = labels.dtype.type(loc)
                except TypeError:
                    # this occurs when loc is a slice (partial string indexing)
                    # but the TypeError raised by searchsorted in this case
                    # is catched in Index._has_valid_type()
                    pass

[jreback]

sure, or just tests if its a scalar to begin with

[codecov-io]

Current coverage is 85.26% (diff: 100%)

Merging #14551 into master will increase coverage by <.01%

@@             master     #14551   diff @@
==========================================
  Files           140        140          
  Lines         50672      50676     +4   
  Methods           0          0          
  Messages          0          0          
  Branches          0          0          
==========================================
+ Hits          43207      43211     +4   
  Misses         7465       7465          
  Partials          0          0          

Powered by Codecov. Last update 60a335e...6447e4c

[jreback]

lgtm. just release note. I think this can only help :>

[jorisvandenbossche]

OK, added whatsnew notice. Will merge, but will open a new issue for follow-up on this (other cases, benchmarks that capture this (at the moment there are none)).