Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Make 'numba' a soft dependency for the Python server API #5169

Merged
merged 2 commits into from
Feb 22, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
Loading