Refactoring and prepare for jarray->np array only

deephaven · Oct 27, 2023 · 710ea00 · 710ea00
1 parent 6bab6c1
commit 710ea00
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 25 deletions.
diff --git a/py/server/deephaven/numpy.py b/py/server/deephaven/numpy.py
@@ -11,7 +11,7 @@
 
 from deephaven import DHError, dtypes, new_table
 from deephaven.column import Column, InputColumn
-from deephaven.dtypes import DType
+from deephaven.dtypes import DType, _PRIMITIVE_DTYPE_NULL_MAP
 from deephaven.table import Table
 
 _JPrimitiveArrayConversionUtility = jpy.get_type("io.deephaven.integrations.common.PrimitiveArrayConversionUtility")
@@ -34,7 +34,7 @@ def column_to_numpy_array(col_def: Column, j_array: jpy.JType) -> np.ndarray:
         raise DHError(e, f"failed to create a numpy array for the column {col_def.name}") from e
 
 
-def _j_array_to_numpy_array(dtype: DType, j_array: jpy.JType) -> np.ndarray:
+def _j_array_to_numpy_array(dtype: DType, j_array: jpy.JType, conv_null: bool = False) -> np.ndarray:
     """ Produces a numpy array from the DType and given Java array."""
     if dtype.is_primitive:
         np_array = np.frombuffer(j_array, dtype.np_type)
@@ -55,6 +55,16 @@ def _j_array_to_numpy_array(dtype: DType, j_array: jpy.JType) -> np.ndarray:
     else:
         np_array = np.array(j_array, np.object_)
 
+    if conv_null:
+        dh_null = _PRIMITIVE_DTYPE_NULL_MAP.get(dtype)
+        if dtype in (dtypes.float32, dtypes.float64):
+            np_array = np.copy(np_array)
+        else:
+            if dtype is dtypes.bool_:  # promote boolean to float64
+                np_array = np.frombuffer(np_array, np.byte)
+            np_array = np_array.astype(np.float64)
+        np_array[np_array == dh_null] = np.nan
+
     return np_array
 
 

diff --git a/py/server/deephaven/pandas.py b/py/server/deephaven/pandas.py
@@ -4,7 +4,7 @@
 
 """ This module supports the conversion between Deephaven tables and pandas DataFrames. """
 from functools import wraps
-from typing import List, Dict
+from typing import List, Dict, Literal, Callable, Any
 
 import jpy
 import numpy as np
@@ -307,7 +307,7 @@ def to_table(df: pd.DataFrame, cols: List[str] = None) -> Table:
 }
 
 
-def _dh_null_conv_params(null_value):
+def _dh_null_conv_params(null_value: Literal[np.nan, pd.NA, None]):
     def decorator(fn):
         """A decorator that replaces Deephaven null values in the parameters with the provided replacement value (
         either numpy.nan or pandas.NA) before calling the decorated function.
@@ -332,32 +332,20 @@ def wrapper(*args, **kwargs):
             for arg, np_dtype_char in zip(args, fn_sig):
                 if np_dtype := _J_ARRAY_NP_TYPE_MAP.get(type(arg)):
                     dtype = dtypes.from_np_dtype(np_dtype)
-                    if null_value is np.nan:
-                        np_array = _j_array_to_numpy_array(dtype, arg)
-                        dh_null = _PRIMITIVE_DTYPE_NULL_MAP.get(dtype)
-                        if dtype in (dtypes.float32, dtypes.float64):
-                            np_array = np.copy(np_array)
-                        else:
-                            if dtype is dtypes.bool_:  # promote boolean to float64
-                                np_array = np.frombuffer(np_array, np.byte)
-                            np_array = np_array.astype(np.float64)
-                        np_array[np_array == dh_null] = null_value
-                        converted_args.append(np_array)
-                    else:  # null_value is pd.NA:
-                        pd_series = _j_array_to_series(dtype, arg, conv_null=True)
-                        converted_args.append(pd_series)
+                    if null_value is pd.NA:
+                        converted_args.append(_j_array_to_series(dtype, arg, conv_null=True))
+                    else:  # np.nan or None
+                        converted_args.append(_j_array_to_numpy_array(dtype, arg, conv_null=bool(null_value)))
                 else:  # scalar type or array types that don't need conversion
                     try:
                         np_dtype = np.dtype(np_dtype_char)
                     except TypeError:
                         converted_args.append(arg)
-                        continue
-                    dtype = dtypes.from_np_dtype(np_dtype)
-                    if dtype is dtypes.bool_:
-                        converted_args.append(null_value if arg is None else arg)
                     else:
-                        dh_null = _PRIMITIVE_DTYPE_NULL_MAP.get(dtype)
-                        if dh_null is not None:
+                        dtype = dtypes.from_np_dtype(np_dtype)
+                        if dtype is dtypes.bool_:
+                            converted_args.append(null_value if arg is None else arg)
+                        elif dh_null := _PRIMITIVE_DTYPE_NULL_MAP.get(dtype):
                             converted_args.append(null_value if arg == dh_null else arg)
                         else:
                             converted_args.append(arg)

diff --git a/py/server/tests/test_pyfunc_param_null.py b/py/server/tests/test_pyfunc_param_null.py
@@ -228,7 +228,7 @@ def custom_func(col1, col2) -> bool:
         @dh_null_to_na
         def custom_func2(col1, col2) -> bool:
             return any(col2.isna())
-    
+
         col1_formula = "Col1 = i % 10"
         col2_formula = f"Col2 = i % 3 == 0? null : true"
         t = empty_table(100).update([col1_formula, col2_formula]).group_by("Col1")