Refactor ak.apply

arkouda/apply.py

@@ -136,16 +136,6 @@ def apply(
     from arkouda.core.client import generic_msg
     from arkouda.numpy.pdarrayclass import create_pdarray
 
-    if getattr(apply, "is_apply_supported", None) is None:
-        res = generic_msg("isPythonModuleSupported")
-        is_supported = parse_single_value(cast(str, res))
-        setattr(apply, "is_apply_supported", is_supported)
-
-    if not getattr(apply, "is_apply_supported", False):
-        raise RuntimeError(
-            "The apply module is not supported by the version of Chapel " + "this server was built with."
-        )
-
     vers_supported = getattr(apply, "is_version_supported", None)
     if vers_supported is None:
         interp_version = f"{sys.version_info.major}.{sys.version_info.minor}"

src/ApplyMsg.chpl

@@ -1,7 +1,7 @@
 module ApplyMsg {
   use NumPyDType;
-  use ArkoudaPythonCompat;
 
+  use Python;
   use BigInteger;
   use BlockDist;
 
@@ -16,6 +16,8 @@ module ApplyMsg {
   use Logging;
   use ServerConfig;
 
+  import this.Base64.b64Decode;
+
   private config const logLevel = ServerConfig.logLevel;
   private config const logChannel = ServerConfig.logChannel;
   const applyLogger = new Logger(logLevel, logChannel);
@@ -26,77 +28,49 @@ module ApplyMsg {
     return interpreters[here.id]!;
   }
 
-  @arkouda.registerCommand
-  proc isPythonModuleSupported(): bool {
-    return pythonModuleSupported;
+  proc pythonVersionString(): string {
+    return CPythonInterface.PY_MAJOR_VERSION:string + "." +
+           CPythonInterface.PY_MINOR_VERSION:string;
   }
 
   @arkouda.registerCommand
   proc initPythonInterpreters() throws {
-    if !pythonModuleSupported {
-      throw new ErrorWithContext(
-        "Python module not supported with this version of Chapel",
-        getLineNumber(),
-        getRoutineName(),
-        getModuleName()
-      );
-    } else {
-      applyLogger.debug(
-        getModuleName(),
-        getRoutineName(),
-        getLineNumber(),
-        "Initializing Python interpreters");
-
-      forall interp in interpreters do
-        interp = new Interpreter();
-    }
+    applyLogger.debug(
+      getModuleName(),
+      getRoutineName(),
+      getLineNumber(),
+      "Initializing Python interpreters");
+
+    forall interp in interpreters do
+      interp = new Interpreter();
   }
 
   @arkouda.registerCommand
   proc isVersionSupported(versionString: string): bool throws {
-    if !pythonModuleSupported {
-      throw new ErrorWithContext(
-        "Python module not supported with this version of Chapel",
-        getLineNumber(),
-        getRoutineName(),
-        getModuleName()
-      );
-      return false;
-    } else {
-      return versionString == ServerConfig.pythonVersion &&
-             versionString == ArkoudaPythonCompat.pythonVersionString();
-    }
+    return versionString == ServerConfig.pythonVersion &&
+            versionString == pythonVersionString();
   }
 
   @arkouda.registerCommand
   proc applyStr(const ref x: [?d] ?t, funcStr: string): [d] t throws {
     var ret = makeDistArray(d, t);
-    if !pythonModuleSupported {
-      throw new ErrorWithContext(
-        "Python module not supported with this version of Chapel",
-        getLineNumber(),
+    coforall l in d.targetLocales() do on l {
+      var func = new Function(getInterpreter(), funcStr);
+      applyLogger.debug(
+        getModuleName(),
         getRoutineName(),
-        getModuleName()
-      );
-    } else {
-      coforall l in d.targetLocales() do on l {
-        var func = new Function(getInterpreter(), funcStr);
-        applyLogger.debug(
-          getModuleName(),
-          getRoutineName(),
-          getLineNumber(),
-          "Loaded function on locale " + l.id:string + ", applying function");
-        for sd in d.localSubdomains() {
-          for i in sd {
-            ret[i] = func(t, x[i]);
-          }
+        getLineNumber(),
+        "Loaded function on locale " + l.id:string + ", applying function");
+      for sd in d.localSubdomains() {
+        for i in sd {
+          ret[i] = func(t, x[i]);
         }
-        applyLogger.debug(
-          getModuleName(),
-          getRoutineName(),
-          getLineNumber(),
-          "Finished applying function on locale " + l.id:string);
       }
+      applyLogger.debug(
+        getModuleName(),
+        getRoutineName(),
+        getLineNumber(),
+        "Finished applying function on locale " + l.id:string);
     }
     return ret;
   }
@@ -112,55 +86,94 @@ module ApplyMsg {
                    type array_dtype_to): MsgTuple throws
       where array_dtype != BigInteger.bigint &&
             array_dtype_to != BigInteger.bigint {
+    const pickleDataStr = msgArgs["pickleData"].toScalar(string);
+    const pickleData: bytes = b64Decode(pickleDataStr);
 
-    if !pythonModuleSupported {
-      throw new ErrorWithContext(
-        "Python module not supported with this version of Chapel",
-        getLineNumber(),
+    var xSym = st[msgArgs["x"]]: SymEntry(array_dtype, array_nd);
+    ref x = xSym.a;
+
+    var d = x.domain;
+    var ret = makeDistArray(d, array_dtype_to);
+
+    coforall l in d.targetLocales() do on l {
+      var func = new Value(getInterpreter(), pickleData);
+      applyLogger.debug(
+        getModuleName(),
         getRoutineName(),
-        getModuleName()
-      );
-    } else {
-      const pickleDataStr = msgArgs["pickleData"].toScalar(string);
-      var pickleData: bytes;
-      {
-        // TODO: this is a big hack,
-        // ideally we properly implement the base64 module
-        // instead of using Python as a polyfill.
-        // even better, we should have a way to just
-        // send bytes from the client to the server
-        var mod = new Module(getInterpreter(), "base64");
-        var b64decode = new Function(mod, "b64decode");
-        pickleData = b64decode(bytes, pickleDataStr);
+        getLineNumber(),
+        "Loaded pickle data on locale " +
+        l.id:string + ", applying function");
+      for sd in d.localSubdomains() {
+        for i in sd {
+          ret[i] = func(array_dtype_to, x[i]);
+        }
       }
+      applyLogger.debug(
+        getModuleName(),
+        getRoutineName(),
+        getLineNumber(),
+        "Finished applying function on locale " + l.id:string);
+    }
+    return st.insert(new shared SymEntry(ret));
+  }
+
+  /*
+    Copied and modified from https://github.com/jabraham17/Base64/blob/v0.1.0/src/Base64.chpl
+  */
+  module Base64 {
+    /*
+      Decode a Base64 string to binary data.
+    */
+    proc b64Decode(x: string): bytes do
+      return b64DecodeImpl(x);
+
+    private param padding = b"=";
+    private proc decodeChar(ch: uint(8),
+                            param plus: bytes, param slash: bytes): int {
+      if ch >= 0x41 && ch <= 0x5A then return (ch - 0x41): int;        // A-Z
+      else if ch >= 0x61 && ch <= 0x7A then return (ch - 0x61 + 26): int; // a-z
+      else if ch >= 0x30 && ch <= 0x39 then return (ch - 0x30 + 52): int; // 0-9
+      else if ch == plus.toByte() then return 62;
+      else if ch == slash.toByte() then return 63;
+      else return -1; // padding or invalid
+    }
 
+    private proc b64DecodeImpl(x: string,
+                               param plus = b"+", param slash = b"/"): bytes {
+      const len = x.size;
+      if len == 0 then return b"";
 
-      var xSym = st[msgArgs["x"]]: SymEntry(array_dtype, array_nd);
-      ref x = xSym.a;
-
-      var d = x.domain;
-      var ret = makeDistArray(d, array_dtype_to);
-
-      coforall l in d.targetLocales() do on l {
-        var func = new Value(getInterpreter(), pickleData);
-        applyLogger.debug(
-          getModuleName(),
-          getRoutineName(),
-          getLineNumber(),
-          "Loaded pickle data on locale " +
-          l.id:string + ", applying function");
-        for sd in d.localSubdomains() {
-          for i in sd {
-            ret[i] = func(array_dtype_to, x[i]);
-          }
+      var result: bytes;
+      var i = 0;
+
+      while i < len {
+        // Skip whitespace/newlines
+        if x.byte[i] == 0x0A || x.byte[i] == 0x0D ||
+          x.byte[i] == 0x20 || x.byte[i] == 0x09 {
+          i += 1;
+          continue;
         }
-        applyLogger.debug(
-          getModuleName(),
-          getRoutineName(),
-          getLineNumber(),
-          "Finished applying function on locale " + l.id:string);
+
+        // Need at least 2 valid chars for a group
+        const a = decodeChar(x.byte[i], plus, slash);
+        const b = if i+1 < len then decodeChar(x.byte[i+1], plus, slash) else 0;
+        const c = if i+2 < len then decodeChar(x.byte[i+2], plus, slash) else 0;
+        const d = if i+3 < len then decodeChar(x.byte[i+3], plus, slash) else 0;
+
+        param mask = 0xFF:uint(8);
+        result.appendByteValues((((a << 2) | (b >> 4)) & mask):uint(8));
+
+        if i+2 < len && x.byte[i+2] != padding.toByte() {
+          result.appendByteValues((((b << 4) | (c >> 2)) & mask):uint(8));
+        }
+        if i+3 < len && x.byte[i+3] != padding.toByte() {
+          result.appendByteValues((((c << 6) | d) & mask):uint(8));
+        }
+
+        i += 4;
       }
-      return st.insert(new shared SymEntry(ret));
+
+      return result;
     }
   }
 }

tests/apply_test.py

@@ -10,15 +10,6 @@
 """
 
 
-def supports_apply():
-    try:
-        return ak.numpy.pdarrayclass.parse_single_value(
-            ak.core.client.generic_msg("isPythonModuleSupported")
-        )
-    except Exception:
-        return False
-
-
 @pytest.mark.requires_chapel_module("ApplyMsg")
 class TestApply:
     def test_apply_docstrings(self):
@@ -30,11 +21,6 @@ def test_apply_docstrings(self):
         )
         assert result.failed == 0, f"Doctest failed: {result.failed} failures"
 
-    @classmethod
-    def setup_class(cls):
-        if not supports_apply():
-            pytest.skip("apply not supported")
-
     @pytest.mark.parametrize("prob_size", pytest.prob_size)
     @pytest.mark.parametrize("dtype", [ak.int64, ak.uint64])
     def test_apply_inc(self, prob_size, dtype):

tests/pandas/extension/series_accessor.py

@@ -15,7 +15,6 @@
 )
 from arkouda.pandas.index import MultiIndex
 from arkouda.pandas.series import Series as ak_Series
-from tests.apply_test import supports_apply
 
 
 def _assert_series_equal_values(s: pd.Series, values):
@@ -424,11 +423,6 @@ def test_argsort_strings(self):
 
 @pytest.mark.requires_chapel_module("ApplyMsg")
 class TestArkoudaSeriesApply:
-    @classmethod
-    def setup_class(cls):
-        if not supports_apply():
-            pytest.skip("apply not supported")
-
     def test_series_accessor_apply_requires_arkouda_backed(self):
         s = pd.Series([1, 2, 3], name="x")
         with pytest.raises(TypeError, match="Series must be Arkouda-backed"):