PERF: optimize algos.take for repeated calls
State: closed

This PR optimizes our internal `take` algorithm, with the following changes:

- Cache the expensive parts that are independent of the actual array values but only dependent on the dtypes / dimension / axis / fill_value (which is often the same for repeated calls), i.e. `maybe_promote` and `_get_take_nd_function`
- The above is done in the generic `take_nd` function as well, but in addition I also added a `take_1d_array` specialized version that assumes the input is already an array and only deals with 1D. This gives a small further speed-up.

An example use case where the effect of this optimization can be noticed nicely is `unstack` in case of the ArrayManager. This is typically a case where we do multiple `take` calls for 1 column (to split it into multiple columns) and thus call it many times for the same dtype/fill_value.

Using the example from the `reshape.py::SimpleReshape` ASV benchmark (this is an example of homogeneous dtype with a single block, so which is fast with the BlockManager using the `_can_fast_transpose` fastpath in `_unstack_frame`):

```python
arrays = [np.arange(100).repeat(100), np.roll(np.tile(np.arange(100), 100), 25)]
index = pd.MultiIndex.from_arrays(arrays)
df = pd.DataFrame(np.random.randn(10000, 4), index=index)
df_am = df._as_manager("array")
```

Master:
```
In [2]: %timeit df.unstack()
1.79 ms ± 108 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [3]: %timeit df_am.unstack()
15.1 ms ± 1.3 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)
```

PR:

```
In [2]: %timeit df.unstack()
1.81 ms ± 160 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [3]: %timeit df_am.unstack()
3.44 ms ± 69.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
```

So for the BlockManager this doesn't matter (it does only a single call to `take_nd` for the column values), but for the ArrayManager version this gives a nice speedup and brings it within 2x of the BlockManager version (for this benchmark).

Labels:
name: Performance
description: Memory or execution speed performance
name: Algos
description: Non-arithmetic algos: value_counts, factorize, sorting, isin, clip, shift, diff

Comments:
As an example benchmark of `unstack` that doesn't use the 1-block fastpath (from `reshape.py::Unstack::time_full_product[category]`):

```
import string

m = 100
n = 50

levels = np.arange(m)
index = pd.MultiIndex.from_product([levels] * 2)
columns = np.arange(n)
indices = np.random.randint(0, 52, size=(m * m, n))
values = np.take(list(string.ascii_letters), indices)
values = [pd.Categorical(v) for v in values.T]
df = pd.DataFrame(values, index, columns)
df_am = df._as_manager("array")
```

Master:

```
In [2]: %timeit df.unstack()
374 ms ± 7.42 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [3]: %timeit df_am.unstack()
306 ms ± 10 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
```

PR:

```
In [2]: %timeit df.unstack()
309 ms ± 21 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [3]: %timeit df_am.unstack()
177 ms ± 27.8 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

So here it improves for BlockManager as well (and in this case ArrayManager is faster, because it has less overhead creating all the blocks for each column)

(specifically for this case of Categorical data, there is a lot of room for improvement by making the categorical constructor faster for such cases where less validation is needed, but that's another topic)
@jorisvandenbossche a general comment. target PRs are much better for this type of refactoring. Sure occasionally bigger ones are needed. But please try to make small incremental changes (I know you probably feel this is a small incremental change), but its not. Meaning break things into multple PRs where you first move things, then change them up. This may seem like more work on your part, and migth be a bit more, but it will make reviews faster and less back and forth and arguing.
I think indeed this is a rather targetted change. A large part of the diff is caused by copying part of a function in helper functions to be able to reuse it (and I commented on those parts to indicate where this happened). 

Sure I can do the move before this PR. 
@jorisvandenbossche the timings in the OP look like this makes the non-ArrayManager code very slightly slower (well within margin of error) for this particular use case.  are there other use cases where this improves perf? 
yea pls remove the alias would help
-> precursor just moving code: https://github.com/pandas-dev/pandas/pull/39728
> the timings in the OP look like this makes the non-ArrayManager code very slightly slower (well within margin of error) for this particular use case. are there other use cases where this improves perf?

It's not slower, but just small variation in timing for something that didn't change noticeably. At least, the precision of the timing is not good enough to know if it's tiny bit slower or tiny bit faster. I could have repeated it to have the difference the other way around ;))

I didn't check where this would improve performance for the BlockManager noticeably. In principle any case where the performance is dominated by the overhead of `take` (so you need a case where `take` is called many times), but don't know from the top of my head what such a case could be. 

I suppose that you could measure the benefit in a single `take` (eg reindex) call as well.
can you add a comment along the lines of `# TODO(AM): this can be removed if we dont end up using ArrayManager`
> can you add a comment along the lines of # TODO(AM): this can be removed if we dont end up using ArrayManager

Where do you want me to add that comment? Specifically on the specialized `take_1d_array`? 
I think the hashing is probably useful in general as well.
> Where do you want me to add that comment?

on any of the new functions where itd be accurate
The sticking point seems to be the complexity introduced by the caching.  two thoughts that might trim that down

1) the hashability rule in maybe_promote is not-quite right.  ATM `maybe_promote(not_object, not_hashable)` raises, but in principle we could have it return `(object, not_hashable)`

2) Since the hashability check can go at the top of maybe_promote, we can then refactor _the rest_ of maybe_promote into  the cached _maybe_promote.
> The sticking point seems to be the complexity introduced by the caching. two thoughts that might trim that down
> 
> 1. the hashability rule in maybe_promote is not-quite right.  ATM `maybe_promote(not_object, not_hashable)` raises, but in principle we could have it return `(object, not_hashable)`
> 2. Since the hashability check can go at the top of maybe_promote, we can then refactor _the rest_ of maybe_promote into  the cached _maybe_promote.

yes I think 2 is the way to go here. I am all for caching, but we have to be very careful to make sure the underlying function that is cached is the simplest possible.
can you merge master
Updated (and also removed the `take_1d` specialization for now, to focus on the caching, since that both gives the most significant improvement as generates the most discussion). 

> 2\. Since the hashability check can go at the top of maybe_promote, we can then refactor _the rest_ of maybe_promote into the cached _maybe_promote.

OK, so that version looks more or less like this (@jbrockmendel did I understand that correctly?)

```python
def maybe_promote_with_cache_version1(dtype: np.dtype, fill_value=np.nan):
    if not is_hashable(fill_value):
        if not is_object_dtype(dtype):
            raise ValueError("fill_value must be a scalar")
        return dtype, fill_value
    return _maybe_promote_cached(dtype, fill_value, type(fill_value))

@functools.lru_cache(maxsize=128)
def _maybe_promote_cached(dtype, fill_value, fill_value_type):
    ... current implementation of maybe_promote ...
```

while the version I have in this PR basically looks like:

```python
def maybe_promote_with_cache_version2(dtype, fill_value):
    try:
        return _maybe_promote_cached(dtype, fill_value, type(fill_value))
    except TypeError:
        return _maybe_promote(dtype, fill_value) 

@functools.lru_cache(maxsize=128)
def _maybe_promote_cached(dtype, fill_value, fill_value_type):
    return _maybe_promote(dtype, fill_value)

def _maybe_promote(dtype, fill_value):
    ... current implementation of maybe_promote ...
```

(and to be clear the `maybe_promote_with_cache_version1/2` would then become the actual `maybe_promote` that is used elsewhere in pandas, just using those names to distinguish and test/time them next to each other)

Timing those options:

```
In [2]: dtype = np.dtype(float)
   ...: fill_value = np.nan

In [3]: %timeit maybe_promote(dtype, fill_value)   # <--- the current, non-cached version
2.36 µs ± 10 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [4]: %timeit maybe_promote_with_cache_version1(dtype, fill_value)
303 ns ± 3.59 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

In [5]: %timeit maybe_promote_with_cache_version2(dtype, fill_value)
201 ns ± 3.91 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
```

So the version with try/except is clearly faster as first version with caching. But of course both are in the nanoseconds, so in absolute numbers the difference is not that big. 

For the example that I use in the top post, `maybe_promote` gets called 400 times, which in this case means 400 x 100ns = 40µs gain by using the second version, while the full unstack operation takes around 3ms. Which of course is only a small precentage.

Personally I don't think both versions are that different in terms of complexity. But so if there is a strong preference for the slightly slower version 1 compared to version 2, I am fine with either (the speed-up on the non-microbenchmark is not very significant). 

@jorisvandenbossche ok to go with your current version, i would just add a comment explaining the try/except (e.g. that non-hashables will except, maybe you did this alreadY).
ill take another look this afternoon
AFAICT there isnt any functional difference between maybe_promote_with_cache_version1 vs maybe_promote_with_cache_version2, so lets go with the faster one.  can you move it to dtypes.cast and we'll get this done
thanks @jorisvandenbossche very nice
@jorisvandenbossche im looking at pandas/tests/indexing/multiindex/test_indexing_slow.py::test_multiindex_get_loc and finding that disabling the lru_cache cuts the runtime in about half (for BM, about 40% for AM).  Can you confirm?

Disabling the cache speeds up the test suite by about 20% for me.
 I take it back, was using an out of date baseline.
Do you have a specific reproducer as a small snippet (outside of pytest)? When running those tests, I do indeed see a speedup when disabling the cache, but running that test includes both pytest overhead as several ways we end up calling take (sort_values, the actual MI getitem, in the assert functions, ..), so might be easier with a narrowed down use case to profile it / understand what is going on.


Trimming down test_multiindex_get_loc to something minimal-ish:

```
import numpy as np
import pandas as pd

m = 50
n = 1000
cols = ["jim", "joe", "jolie", "joline", "jolia"]

vals = [
    np.random.randint(0, 10, n),
    np.random.choice(list("abcdefghij"), n),
    np.random.choice(pd.date_range("20141009", periods=10).tolist(), n),
    np.random.choice(list("ZYXWVUTSRQ"), n),
    np.random.randn(n),
]
vals = list(map(tuple, zip(*vals)))

df = pd.DataFrame(vals, columns=cols)
mi = df.set_index(cols[:-1])

key = (8, 'h', pd.Timestamp('2014-10-17 00:00:00'), 'S')
i = 0
k = key[i]

mask = df.iloc[:, i] == k

%timeit right = df[mask]
1.57 ms ± 53.9 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)  # <- main
174 µs ± 4.26 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)  # <- cache disabled
```

Removing the date_range from `vals` (and an entry from `cols`) brings the two back to having the cache case outperforming by a few percent.

First guess is that for some reason dt64NaT is causing a cache miss, but calling _maybe_promote_cached with (np.dtype("M8[ns]"), np.datetime64("NaT", "ns), np.datetime64) isn't showing a miss.
So it seems that it is "taking" the datetimelike values then, so further trimming it down to just that:

```
n = 1000
vals = np.random.choice(pd.date_range("20141009", periods=10).tolist(), n)
ser = pd.Series(vals)
mask = np.random.randint(0, 10, n) == 0

%timeit ser[mask]

df = pd.DataFrame({"col": vals})

%timeit df[mask]
```

Some observations:

- The slowdown only occurs in the dataframe case, not for series. But that's simply because the series case's code path doesn't pass by those cache functions
- It's due to `_maybe_promote_cached`, not the `_get_take_nd_function_cached`
- I checked the time of hashing both `dtype = np.dtype('<M8[ns]')` and `fill_value = np.datetime64('NaT')`, but nothing seems concerning about this
- Also timing explicitly the `_maybe_promote_(cached)` doesn't seem to show anything:

```
In [1]: dtype = np.dtype('<M8[ns]')
   ...: fill_value = np.datetime64('NaT')

# version without cache
In [2]: %timeit pd.core.dtypes.cast._maybe_promote(dtype, fill_value)
2.8 µs ± 49.9 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

In [3]: %timeit pd.core.dtypes.cast.maybe_promote(dtype, fill_value)
468 ns ± 7.62 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)
```


Thanks for taking a look.  I get similar numbers when profiling just the _?maybe_promote functions.

Speculation: each time we call maybe_promote with (np.dtype("M8[ns]"), np.datetime64("NaT", "ns")), we are actually passing a new dt64 object.  When doing the dict lookup it correctly gets hash-equality and then checks for list-equality, which then shows up as False.  So we get cache misses.

Let's check this out...  At the module-level in dtypes.cast I define `dt64nat = np.datetime64("NaT")` and in maybe_promote I do

```
    if isinstance(fill_value, np.datetime64) and isna(fill_value):
        fill_value = dt64nat
```

Then using the example from https://github.com/pandas-dev/pandas/pull/39692#issuecomment-1439203636

```
In [2]: %timeit right = df[mask]
182 µs ± 7.12 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
```

about 8 microseconds slower than the disabling-the-cache measurement, which seems reasonable for the extra check added.  Calling this a win.

Presumably we'd see the same issue with td64nat or if any case we passed a NaN that wasn't specifically `np.nan` (which is probably rare).