Fix imports; change binary from hex to base64

Author: kesmit13

singlestoredb/functions/dtypes.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+import base64
 import datetime
 import decimal
 import re
@@ -106,7 +107,7 @@ def bytestr(x: Any) -> Optional[bytes]:
         return x
     if isinstance(x, bytes):
         return x
-    return bytes.fromhex(x)
+    return base64.b64decode(x)
 
 
 PYTHON_CONVERTERS = {

singlestoredb/functions/ext/json.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+import base64
 import json
 from typing import Any
 from typing import List
@@ -35,7 +36,7 @@ class JSONEncoder(json.JSONEncoder):
 
     def default(self, obj: Any) -> Any:
         if isinstance(obj, bytes):
-            return obj.hex()
+            return base64.b64encode(obj).decode('utf-8')
         return json.JSONEncoder.default(self, obj)
 
 
@@ -130,6 +131,8 @@ def load_pandas(
     Tuple[pd.Series[int], List[pd.Series[Any]]
 
     '''
+    import numpy as np
+    import pandas as pd
     row_ids, cols = _load_vectors(colspec, data)
     index = pd.Series(row_ids, dtype=np.longlong)
     return index, \
@@ -164,6 +167,7 @@ def load_polars(
     Tuple[polars.Series[int], List[polars.Series[Any]]
 
     '''
+    import polars as pl
     row_ids, cols = _load_vectors(colspec, data)
     return pl.Series(None, row_ids, dtype=pl.Int64), \
         [
@@ -194,6 +198,7 @@ def load_numpy(
     Tuple[np.ndarray[int], List[np.ndarray[Any]]
 
     '''
+    import numpy as np
     row_ids, cols = _load_vectors(colspec, data)
     return np.asarray(row_ids, dtype=np.longlong), \
         [
@@ -224,6 +229,7 @@ def load_arrow(
     Tuple[pyarrow.Array[int], List[pyarrow.Array[Any]]
 
     '''
+    import pyarrow as pa
     row_ids, cols = _load_vectors(colspec, data)
     return pa.array(row_ids, type=pa.int64()), \
         [

singlestoredb/functions/ext/rowdat_1.py

@@ -35,7 +35,7 @@
     except ImportError:
         pass
     try:
-        import pyarrow.compute as pc
+        import pyarrow.compute as pc  # noqa: F401
     except ImportError:
         pass
 
@@ -205,6 +205,7 @@ def _load_pandas(
     Tuple[pd.Series[int], List[Tuple[pd.Series[Any], pd.Series[bool]]]]
 
     '''
+    import numpy as np
     import pandas as pd
 
     row_ids, cols = _load_vectors(colspec, data)
@@ -558,6 +559,7 @@ def _load_pandas_accel(
     if not has_accel:
         raise RuntimeError('could not load SingleStoreDB extension')
 
+    import numpy as np
     import pandas as pd
 
     numpy_ids, numpy_cols = _singlestoredb_accel.load_rowdat_1_numpy(colspec, data)
@@ -663,8 +665,11 @@ def _create_arrow_mask(
     data: 'pa.Array[Any]',
     mask: 'pa.Array[pa.bool_]',
 ) -> 'pa.Array[pa.bool_]':
+    import pyarrow.compute as pc  # noqa: F811
+
     if mask is None:
         return data.is_null().to_numpy(zero_copy_only=False)
+
     return pc.or_(data.is_null(), mask.is_null()).to_numpy(zero_copy_only=False)
 
 

singlestoredb/functions/utils.py

@@ -6,12 +6,6 @@
 from typing import Any
 from typing import Dict
 
-try:
-    import numpy as np
-    has_numpy = True
-except ImportError:
-    has_numpy = False
-
 
 if sys.version_info >= (3, 10):
     _UNION_TYPES = {typing.Union, types.UnionType}
@@ -36,29 +30,51 @@ def get_annotations(obj: Any) -> Dict[str, Any]:
     return getattr(obj, '__annotations__', {})
 
 
+def get_module(obj: Any) -> str:
+    """Get the module of an object."""
+    module = getattr(obj, '__module__', '').split('.')
+    if module:
+        return module[0]
+    return ''
+
+
+def get_type_name(obj: Any) -> str:
+    """Get the type name of an object."""
+    if hasattr(obj, '__name__'):
+        return obj.__name__
+    if hasattr(obj, '__class__'):
+        return obj.__class__.__name__
+    return ''
+
+
 def is_numpy(obj: Any) -> bool:
     """Check if an object is a numpy array."""
-    if is_union(obj):
-        obj = typing.get_args(obj)[0]
-    if not has_numpy:
-        return False
     if inspect.isclass(obj):
-        return obj is np.ndarray
-    if typing.get_origin(obj) is np.ndarray:
-        return True
-    return isinstance(obj, np.ndarray)
+        if get_module(obj) == 'numpy':
+            return get_type_name(obj) == 'ndarray'
+
+    origin = typing.get_origin(obj)
+    if get_module(origin) == 'numpy':
+        if get_type_name(origin) == 'ndarray':
+            return True
+
+    dtype = type(obj)
+    if get_module(dtype) == 'numpy':
+        return get_type_name(dtype) == 'ndarray'
+
+    return False
 
 
 def is_dataframe(obj: Any) -> bool:
     """Check if an object is a DataFrame."""
     # Cheating here a bit so we don't have to import pandas / polars / pyarrow:
     # unless we absolutely need to
-    if getattr(obj, '__module__', '').startswith('pandas.'):
-        return getattr(obj, '__name__', '') == 'DataFrame'
-    if getattr(obj, '__module__', '').startswith('polars.'):
-        return getattr(obj, '__name__', '') == 'DataFrame'
-    if getattr(obj, '__module__', '').startswith('pyarrow.'):
-        return getattr(obj, '__name__', '') == 'Table'
+    if get_module(obj) == 'pandas':
+        return get_type_name(obj) == 'DataFrame'
+    if get_module(obj) == 'polars':
+        return get_type_name(obj) == 'DataFrame'
+    if get_module(obj) == 'pyarrow':
+        return get_type_name(obj) == 'Table'
     return False
 
 
@@ -74,13 +90,13 @@ def get_data_format(obj: Any) -> str:
     """Return the data format of the DataFrame / Table / vector."""
     # Cheating here a bit so we don't have to import pandas / polars / pyarrow
     # unless we absolutely need to
-    if getattr(obj, '__module__', '').startswith('pandas.'):
+    if get_module(obj) == 'pandas':
         return 'pandas'
-    if getattr(obj, '__module__', '').startswith('polars.'):
+    if get_module(obj) == 'polars':
         return 'polars'
-    if getattr(obj, '__module__', '').startswith('pyarrow.'):
+    if get_module(obj) == 'pyarrow':
         return 'arrow'
-    if getattr(obj, '__module__', '').startswith('numpy.'):
+    if get_module(obj) == 'numpy':
         return 'numpy'
     if isinstance(obj, list):
         return 'list'
@@ -92,8 +108,8 @@ def is_pandas_series(obj: Any) -> bool:
     if is_union(obj):
         obj = typing.get_args(obj)[0]
     return (
-        getattr(obj, '__module__', '').startswith('pandas.') and
-        getattr(obj, '__name__', '') == 'Series'
+        get_module(obj) == 'pandas' and
+        get_type_name(obj) == 'Series'
     )
 
 
@@ -102,8 +118,8 @@ def is_polars_series(obj: Any) -> bool:
     if is_union(obj):
         obj = typing.get_args(obj)[0]
     return (
-        getattr(obj, '__module__', '').startswith('polars.') and
-        getattr(obj, '__name__', '') == 'Series'
+        get_module(obj) == 'polars' and
+        get_type_name(obj) == 'Series'
     )
 
 
@@ -112,8 +128,8 @@ def is_pyarrow_array(obj: Any) -> bool:
     if is_union(obj):
         obj = typing.get_args(obj)[0]
     return (
-        getattr(obj, '__module__', '').startswith('pyarrow.') and
-        getattr(obj, '__name__', '') == 'Array'
+        get_module(obj) == 'pyarrow' and
+        get_type_name(obj) == 'Array'
     )
 
 
@@ -147,6 +163,6 @@ def is_pydantic(obj: Any) -> bool:
     # the class is a subclass
     return bool([
         x for x in inspect.getmro(obj)
-        if x.__module__.startswith('pydantic.')
-        and x.__name__ == 'BaseModel'
+        if get_module(x) == 'pydantic'
+        and get_type_name(x) == 'BaseModel'
     ])