ENH: Add masked engine #49420

phofl · 2022-10-31T18:25:55Z

closes #xxxx (Replace xxxx with the GitHub issue number)
Tests added and passed if fixing a bug or adding a new feature
All code checks passed.
Added type annotations to new arguments/methods/functions.
Added an entry in the latest doc/source/whatsnew/vX.X.X.rst file if fixing a bug or adding a new feature.

       before           after         ratio
     [626b6510]       [efde9dd0]
-        270±10μs          148±3μs     0.55  indexing.NumericMaskedIndexing.time_get_indexer('UInt64', True)
-         210±5μs          106±1μs     0.51  indexing.NumericMaskedIndexing.time_get_indexer('Int64', True)
-        306±50μs          150±3μs     0.49  indexing.NumericMaskedIndexing.time_get_indexer('UInt64', False)
-       340±100μs          158±3μs     0.46  indexing.NumericMaskedIndexing.time_get_indexer_dups('Float64', False)
-        237±20μs        107±0.6μs     0.45  indexing.NumericMaskedIndexing.time_get_indexer('Int64', False)
-        360±70μs          157±2μs     0.44  indexing.NumericMaskedIndexing.time_get_indexer_dups('Float64', True)
-       369±100μs          155±3μs     0.42  indexing.NumericMaskedIndexing.time_get_indexer('Float64', False)
-        330±80μs          137±1μs     0.42  indexing.NumericMaskedIndexing.time_get_indexer_dups('UInt64', True)
-        335±50μs          140±2μs     0.42  indexing.NumericMaskedIndexing.time_get_indexer_dups('UInt64', False)
-       373±200μs          155±3μs     0.42  indexing.NumericMaskedIndexing.time_get_indexer('Float64', True)
-        270±30μs          106±1μs     0.39  indexing.NumericMaskedIndexing.time_get_indexer_dups('Int64', False)
-        310±60μs          110±1μs     0.36  indexing.NumericMaskedIndexing.time_get_indexer_dups('Int64', True)

Lets see if ci passes

jbrockmendel · 2022-11-02T03:39:55Z

asv_bench/benchmarks/indexing_engines.py

@@ -9,6 +9,11 @@

 from pandas._libs import index as libindex

+from pandas import (
+    NA,
+    Series,


This file shouldn’t depend on series (see docstring)

We still need to construct the data that is passed to the Engine constructor in some way. For the other engines that is a plain numpy array (so not depending on anything outside of pd._libs), but for those new engines for masked arrays that are the data and mask numpy arrays. While it is possible to construct those manually, it is easier to do that through pd.Series(..) (or pd.array(..) would be sufficient), I think? (and this is only for the setup, not for the actual benchmark)

Thx missed that docstring

We don't need the actual array since we are passing the underlying NumPy arrays anyway, we can just construct the mask and the data separately

jbrockmendel · 2022-11-07T22:10:10Z

pandas/_libs/index.pyx

@@ -158,6 +168,8 @@ cdef class IndexEngine:
            return self._get_loc_duplicates(val)

        try:
+            if self.mask is not None and val is C_NA:
+                return self.mapping.get_na()


can this go outside the try/except?

pandas/_libs/index.pyx

jbrockmendel · 2022-11-07T22:12:47Z

pandas/core/indexes/base.py

-        ):
-            return libindex.ExtensionEngine(target_values)
+        if isinstance(target_values, ExtensionArray):
+            is_masked = hasattr(target_values, "_mask")


can you check directly for e.g. BaseMaskedArray here?

Yes, have to import BaseMaskedArray locally though

pandas/core/indexes/base.py

# Conflicts: # asv_bench/benchmarks/indexing.py # pandas/tests/indexes/numeric/test_indexing.py

pandas/_libs/hashtable.pyi

pandas/_libs/hashtable_class_helper.pxi.in

pandas/_libs/index.pyx

jbrockmendel · 2022-11-30T19:52:37Z

pandas/core/indexes/base.py

-        ):
-            return libindex.ExtensionEngine(target_values)
+        if isinstance(target_values, ExtensionArray):
+            from pandas.core.arrays import BaseMaskedArray


cant do this up top?

No, circular import...

hmm that sucks. not for this PR but i bet either Categorical or IntervalArray are importing something from indexes and it would be reasonable to make the runtime imports happen there

Yep sounds sensible

pandas/core/indexes/base.py

jbrockmendel · 2022-11-30T19:55:01Z

pandas/core/indexes/base.py

+        if isinstance(vals, StringArray):
+            # GH#45652 much more performant than ExtensionEngine
+            return vals._ndarray
+        if type(self) is Index and isinstance(self._values, ExtensionArray):
            return self._values.astype(object)


IIRC _can_use_libjoin will make it so we don't go down this path

Yes, but we can not use this in one case, because the no use lib join path does not work there for non interval, so we have to go through object and end up here

phofl · 2022-12-19T22:17:04Z

This is less ugly now :)

jbrockmendel · 2023-01-01T21:52:00Z

pandas/core/indexes/base.py

@@ -4762,8 +4779,9 @@ def _get_engine_target(self) -> ArrayLike:
            # GH#45652 much more performant than ExtensionEngine
            return vals._ndarray
        if type(self) is Index and isinstance(self._values, ExtensionArray):
-            # TODO(ExtensionIndex): remove special-case, just use self._values
-            return self._values.astype(object)
+            if not isinstance(self._values, BaseMaskedArray):


an this be combined with the previous condition?

Done, found this a bit easier to read, but no strong opinion

phofl · 2023-01-16T22:27:58Z

cc @jbrockmendel Could you have another look?

jbrockmendel · 2023-01-23T22:39:58Z

asv_bench/benchmarks/indexing.py

+        indices = {
+            True: Index(range(N), dtype=dtype),
+            False: Index(
+                list(range(50)) + [54, 53, 52, 51] + list(range(55, N - 1)), dtype=dtype


NBD but i think if you do this at the class level outside of setup it might only construct once. might add up

Do we run into risks of caching stuff on the Index level that could impact performance?

its possible. could just create the arrays instead of the Index objects

Good point, moved the list construction outside. We should probably think about going over the asvs, we are doing this in a lot of places.

jbrockmendel · 2023-01-23T22:41:05Z

asv_bench/benchmarks/indexing_engines.py

@@ -9,6 +9,8 @@

 from pandas._libs import index as libindex

+from pandas.core.arrays import BaseMaskedArray


this renders the file-level docstring not-quite accurate. It shouldn't affect the main point of the docstring, but might be worth noting.

jbrockmendel

LGTM thanks for being patient

…ed_engine

mroeschke

Thanks. Could use a whatsnew note describing the performance improvement of indexing with masked dtypes

phofl · 2023-01-26T00:35:42Z

Added

mroeschke · 2023-01-26T19:36:30Z

Thanks @phofl

jbrockmendel · 2023-01-26T19:44:28Z

Nice! Way to press through

phofl · 2023-01-27T00:38:20Z

Thanks!

I'll work on deduplicating some parts of the code soonish (as discussed sometime during the review)

phofl added 5 commits October 31, 2022 16:54

Implement masked engine

6f5a8ea

Add tests

a051815

Fix asv

3513430

Fix mypy

2246123

Add test

1d35e2e

phofl added Enhancement Indexing Related to indexing on series/frames, not to indexes themselves labels Oct 31, 2022

phofl added 4 commits October 31, 2022 20:22

Fix error

797269a

Fix windows builds

32e6db4

Merge remote-tracking branch 'upstream/main' into enh_masked_engine

9380873

Fix typing

36a0d29

jbrockmendel reviewed Nov 2, 2022

View reviewed changes

Use np arrays

ddcdb13

jbrockmendel reviewed Nov 7, 2022

View reviewed changes

pandas/_libs/index.pyx Outdated Show resolved Hide resolved

jbrockmendel reviewed Nov 7, 2022

View reviewed changes

pandas/core/indexes/base.py Outdated Show resolved Hide resolved

phofl added 5 commits November 10, 2022 21:52

Adress review

e29a970

Merge remote-tracking branch 'upstream/main' into enh_masked_engine

d258bc3

Merge remote-tracking branch 'upstream/main' into enh_masked_engine

6fac3b8

# Conflicts: # asv_bench/benchmarks/indexing.py # pandas/tests/indexes/numeric/test_indexing.py

Adapt to join difference

aec65a7

Merge remote-tracking branch 'upstream/main' into enh_masked_engine

f0e749a