Skip to content

Commit

Permalink
Make 'numba' a soft dependency for the Python server API (#5169)
Browse files Browse the repository at this point in the history
* Make 'numba' a soft_dependency

* Apply suggestions from code review

Co-authored-by: Chip Kent <5250374+chipkent@users.noreply.github.com>

---------

Co-authored-by: Chip Kent <5250374+chipkent@users.noreply.github.com>
  • Loading branch information
2 people authored and stanbrub committed Feb 28, 2024
1 parent 3fb6fd8 commit 21e39dc
Show file tree
Hide file tree
Showing 2 changed files with 76 additions and 49 deletions.
21 changes: 21 additions & 0 deletions py/server/deephaven/_dep.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
#
# Copyright (c) 2016-2023 Deephaven Data Labs and Patent Pending
#
"""This module is an internal module to support dependency management"""

import importlib
from types import ModuleType
from typing import Optional


def soft_dependency(module_name: str) -> Optional[ModuleType]:
"""Attempt to import a module and return it if it exists, otherwise return None.
Args:
module_name (str): the name of the module to import
Returns:
the module if it exists, otherwise None
"""
try:
return importlib.import_module(module_name)
except ImportError:
return None
104 changes: 55 additions & 49 deletions py/server/deephaven/_udf.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,10 @@
from functools import wraps
from typing import Callable, List, Any, Union, Tuple, _GenericAlias

import numba
from deephaven._dep import soft_dependency

numba = soft_dependency("numba")

import numpy
import numpy as np

Expand Down Expand Up @@ -162,56 +165,57 @@ def _parse_return_annotation(annotation: Any) -> _ParsedReturnAnnotation:
return pra


def _parse_numba_signature(fn: Union[numba.np.ufunc.gufunc.GUFunc, numba.np.ufunc.dufunc.DUFunc]) -> _ParsedSignature:
""" Parse a numba function's signature"""
sigs = fn.types # in the format of ll->l, ff->f,dd->d,OO->O, etc.
if sigs:
p_sig = _ParsedSignature(fn)

# for now, we only support one signature for a numba function because the query engine is not ready to handle
# multiple signatures for vectorization https://github.com/deephaven/deephaven-core/issues/4762
sig = sigs[0]
params, rt_char = sig.split("->")

p_sig.params = []
p_sig.ret_annotation = _ParsedReturnAnnotation()
p_sig.ret_annotation.encoded_type = rt_char

if isinstance(fn, numba.np.ufunc.dufunc.DUFunc):
for p in params:
pa = _ParsedParamAnnotation()
pa.encoded_types.add(p)
if p in _NUMPY_INT_TYPE_CODES:
pa.int_char = p
if p in _NUMPY_FLOATING_TYPE_CODES:
pa.floating_char = p
p_sig.params.append(pa)
else: # GUFunc
# An example: @guvectorize([(int64[:], int64[:], int64[:])], "(m),(n)->(n)"
input_output_decl = fn.signature # "(m),(n)->(n)" in the above example
input_decl, output_decl = input_output_decl.split("->")
# remove the parentheses so that empty string indicates no array, non-empty string indicates array
input_decl = re.sub("[()]", "", input_decl).split(",")
output_decl = re.sub("[()]", "", output_decl)

for p, d in zip(params, input_decl):
pa = _ParsedParamAnnotation()
if d:
pa.encoded_types.add("[" + p)
pa.has_array = True
else:
if numba:
def _parse_numba_signature(fn: Union[numba.np.ufunc.gufunc.GUFunc, numba.np.ufunc.dufunc.DUFunc]) -> _ParsedSignature:
""" Parse a numba function's signature"""
sigs = fn.types # in the format of ll->l, ff->f,dd->d,OO->O, etc.
if sigs:
p_sig = _ParsedSignature(fn)

# for now, we only support one signature for a numba function because the query engine is not ready to handle
# multiple signatures for vectorization https://github.com/deephaven/deephaven-core/issues/4762
sig = sigs[0]
params, rt_char = sig.split("->")

p_sig.params = []
p_sig.ret_annotation = _ParsedReturnAnnotation()
p_sig.ret_annotation.encoded_type = rt_char

if isinstance(fn, numba.np.ufunc.dufunc.DUFunc):
for p in params:
pa = _ParsedParamAnnotation()
pa.encoded_types.add(p)
if p in _NUMPY_INT_TYPE_CODES:
pa.int_char = p
if p in _NUMPY_FLOATING_TYPE_CODES:
pa.floating_char = p
p_sig.params.append(pa)

if output_decl:
p_sig.ret_annotation.has_array = True
return p_sig
else:
raise DHError(message=f"numba decorated functions must have an explicitly defined signature: {fn}")
p_sig.params.append(pa)
else: # GUFunc
# An example: @guvectorize([(int64[:], int64[:], int64[:])], "(m),(n)->(n)"
input_output_decl = fn.signature # "(m),(n)->(n)" in the above example
input_decl, output_decl = input_output_decl.split("->")
# remove the parentheses so that empty string indicates no array, non-empty string indicates array
input_decl = re.sub("[()]", "", input_decl).split(",")
output_decl = re.sub("[()]", "", output_decl)

for p, d in zip(params, input_decl):
pa = _ParsedParamAnnotation()
if d:
pa.encoded_types.add("[" + p)
pa.has_array = True
else:
pa.encoded_types.add(p)
if p in _NUMPY_INT_TYPE_CODES:
pa.int_char = p
if p in _NUMPY_FLOATING_TYPE_CODES:
pa.floating_char = p
p_sig.params.append(pa)

if output_decl:
p_sig.ret_annotation.has_array = True
return p_sig
else:
raise DHError(message=f"numba decorated functions must have an explicitly defined signature: {fn}")


def _parse_np_ufunc_signature(fn: numpy.ufunc) -> _ParsedSignature:
Expand All @@ -232,9 +236,11 @@ def _parse_np_ufunc_signature(fn: numpy.ufunc) -> _ParsedSignature:
def _parse_signature(fn: Callable) -> _ParsedSignature:
""" Parse the signature of a function """

if isinstance(fn, (numba.np.ufunc.gufunc.GUFunc, numba.np.ufunc.dufunc.DUFunc)):
return _parse_numba_signature(fn)
elif isinstance(fn, numpy.ufunc):
if numba:
if isinstance(fn, (numba.np.ufunc.gufunc.GUFunc, numba.np.ufunc.dufunc.DUFunc)):
return _parse_numba_signature(fn)

if isinstance(fn, numpy.ufunc):
return _parse_np_ufunc_signature(fn)
else:
p_sig = _ParsedSignature(fn=fn)
Expand Down

0 comments on commit 21e39dc

Please sign in to comment.