[FEA] Improve dtype dispatch patterns in Python
- Dominant language
- C++
- Stars
- 9.8k
- Forks
- 1.1k
- Avg merge
- 3d 6m
- Merged PRs (30d)
- 278
Description
**Is your feature request related to a problem? Please describe.**
There are numerous places in cudf that require polymorphic behavior depending on the dtype of a parameter. Two of the most prominent are `ColumnBase.astype` and `build_column`, both of which essentially boil down to a switch statement based on the dtype.
There are numerous issues with this implementation. Two of the most obvious ones are:
- Performance: See #12494
- It violates the open-closed principle. There is no way to register new types. For `ColumnBase.astype`, this is managed by having the various `as_*_dtype` methods defined in `ColumnBase`, but there is no way to add a new one in general. The only real option would be to monkey-patch the method (same for `build_column`).
- It violates the single responsibility principle. The logic for a new data type's column leaks into the parent ColumnBase's module as well as into helper functions.
**Describe the solution you'd like**
The ideal solution here would be something like `functools.singledispatch` where we register a new override for each dtype and input types are cached to avoid needing to do the expensive `is_*_dtype` calls each time. Unfortunately, this approach is not directly viable because the caching of `singledispatch` is based on the class of the input value, and dtypes cannot be differentiated in this way. For instance, any dtype may be specified as a string.
We should implement a TypeDispatcher (name modeled after libcudf's) that can handle the required type-based single dispatch. The dispatcher could cache the determination of the dtype and run the appropriate method. Unlike `singledispatch`, the dispatch would cache values rather than types. Here is a quick example of what this could look:
```python
from functools import lru_cache
class TypeDispatcher:
def __init__(self, func):
self._default_func = func
self._funcs = []
self._preds = []
def register(self, predicate):
# Not trying to be thread-safe here.
self._preds.append(predicate)
def register_typ(func):
self._funcs.append(func)
return func
return register_typ
@lru_cache
def _get_func(self, obj):
for pred, func in zip(self._preds, self._funcs):
if pred(obj):
return func
return self._default_func
# May want to do some functools.wraps magic to match the signature.
def __call__(self, x, *args, **kwargs):
return self._get_func(x)(x, *args, **kwargs)
@TypeDispatcher
def f(x):
raise ValueError("Unsupported type")
@f.register(lambda x: isinstance(x, int))
def _(x):
return x * 2
@f.register(lambda x: isinstance(x, str))
def _(x):
return x * 3
```
The crucial benefit of this approach (even beyond the performance benefits of the cache) is the ability of external code to register new overloads. Creating a new type of column would no longer require modification of the base column, making our hierarchy much more extensible and removing one major roadblock to supporting alternate column types. Additionally, from an organizational standpoint, it makes the code much easier to follow when all the code relevant to a new column is in one place.
**Describe alternatives you've considered**
Depending on how many different places implement this sort of dispatch, it could be beneficial to decouple the predicates from the functions so that the predicates can be reused, since by assumption they will be identical across all relevant functions. In other words, all instances of `TypeDispatcher` could share a global registry mapping cached objects to their types, but each instance could store a separate mapping from types to functions. That would allow a centralized registration of predicates and even more cache benefits. However, it would require a couple of extra things that may not be worthwhile:
1. Standardizing the way in which the types are keyed so that registration could be done identically everywhere e.g. `@f.register("int_like")`
2. Implementing separate APIs for registering new types and predicates to the class vs registering a new overload for a specific `TypeDispatcher` instance.
**Additional context**
For primitive predicates like the one in my example above, I would expect that the overhead of `lru_cache` and dictionary lookups would be slower than the naive if-else cascade. My assumption is that for the more complex `is_*_dtype` predicates used for dtype-dispatched functions this is not the case.
It's also worth noting that this value-based approach has pitfalls in general around what is considered equal, but those should not bite us when used for dtypes. For instance, `isinstance(1, int)` is True but so is `isinstance(1, float)`. That shouldn't cause us any particular issues though since we are working with dtypes and there are no such overlaps.
This approach may need some modification to work with class methods instead of free functions.
The [multimethod package](https://multimethod.readthedocs.io/en/latest/readme.html#overload) supports this sort of predicate-based overloading (in addition to type-based dispatch), but I think our needs are narrow enough that we are better off at least prototyping this ourselves before reaching to add another dependency.
Contributor guide
Assessment
This issue has not been assessed yet.