ARROW-16211: [C++][Python] Unregister compute functions

python/pyarrow/_compute.pxd

@@ -27,6 +27,10 @@ cdef class ScalarUdfContext(_Weakrefable):
 
     cdef void init(self, const CScalarUdfContext& c_context)
 
+cdef class FunctionRegistry(_Weakrefable):
+    cdef:
+        CFunctionRegistry* registry
+
 cdef class FunctionOptions(_Weakrefable):
     cdef:
         shared_ptr[CFunctionOptions] wrapped

python/pyarrow/_compute.pyx

@@ -477,10 +477,20 @@ cdef _pack_compute_args(object values, vector[CDatum]* out):
 
 
 cdef class FunctionRegistry(_Weakrefable):
-    cdef CFunctionRegistry* registry
 
-    def __init__(self):
-        self.registry = GetFunctionRegistry()
+    def __init__(self, parent=None):
+        cdef:
+            FunctionRegistry new_rg
+            FunctionRegistry parent_rg
+        if parent is None:
+            self.registry = GetFunctionRegistry()
+        else:
+            parent_rg = <FunctionRegistry>(parent)
+            self.registry = CFunctionRegistry.Make(
+                parent_rg.registry).release()
+
+    def __dealloc__(self):
+        self.clear()
 
     def list_functions(self):
         """
@@ -505,6 +515,8 @@ cdef class FunctionRegistry(_Weakrefable):
             func = GetResultValue(self.registry.GetFunction(c_name))
         return wrap_function(func)
 
+    def clear(self):
+        del self.registry
 
 cdef FunctionRegistry _global_func_registry = FunctionRegistry()
 
@@ -2515,7 +2527,7 @@ def _get_scalar_udf_context(memory_pool, batch_length):
 
 
 def register_scalar_function(func, function_name, function_doc, in_types,
-                             out_type):
+                             out_type, registry=function_registry()):
     """
     Register a user-defined scalar function.
 
@@ -2556,6 +2568,12 @@ def register_scalar_function(func, function_name, function_doc, in_types,
         arity.
     out_type : DataType
         Output type of the function.
+    registry: FunctionRegistry
+        FunctionRegistry with which the function will be registered.
+        The default value is set to the global function registry.
+        This is an optional feature to allow grouping functions into
+        different registeries to enable removing functions if they 
+        are not intended to be used further.
 
     Examples
     --------
@@ -2593,6 +2611,8 @@ def register_scalar_function(func, function_name, function_doc, in_types,
         PyObject* c_function
         shared_ptr[CDataType] c_out_type
         CScalarUdfOptions c_options
+        CFunctionRegistry* c_registry
+        FunctionRegistry func_registry
 
     if callable(func):
         c_function = <PyObject*>func
@@ -2628,11 +2648,14 @@ def register_scalar_function(func, function_name, function_doc, in_types,
 
     c_out_type = pyarrow_unwrap_data_type(ensure_type(out_type))
 
+    func_registry = <FunctionRegistry>(registry)
+
     c_options.func_name = c_func_name
     c_options.arity = c_arity
     c_options.func_doc = c_func_doc
     c_options.input_types = c_in_types
     c_options.output_type = c_out_type
+    c_options.registry = func_registry.registry
 
     check_status(RegisterScalarFunction(c_function,
                                         <function[CallbackUdf]> &_scalar_udf_callback, c_options))

python/pyarrow/includes/libarrow.pxd

@@ -1853,6 +1853,7 @@ cdef extern from "arrow/compute/api.h" namespace "arrow::compute" nogil:
         CExecContext()
         CExecContext(CMemoryPool* pool)
         CExecContext(CMemoryPool* pool, CExecutor* exc)
+        CExecContext(CMemoryPool* pool, CExecutor* exc, CFunctionRegistry* rgr)
 
         CMemoryPool* memory_pool() const
         CExecutor* executor()
@@ -1956,6 +1957,9 @@ cdef extern from "arrow/compute/api.h" namespace "arrow::compute" nogil:
         vector[c_string] GetFunctionNames() const
         int num_functions() const
 
+        @staticmethod
+        unique_ptr[CFunctionRegistry] Make(CFunctionRegistry* parent)
+
     CFunctionRegistry* GetFunctionRegistry()
 
     cdef cppclass CElementWiseAggregateOptions \
@@ -2762,13 +2766,15 @@ cdef extern from "arrow/python/udf.h" namespace "arrow::py":
     cdef cppclass CScalarUdfContext" arrow::py::ScalarUdfContext":
         CMemoryPool *pool
         int64_t batch_length
+        CFunctionRegistry *registry
 
     cdef cppclass CScalarUdfOptions" arrow::py::ScalarUdfOptions":
         c_string func_name
         CArity arity
         CFunctionDoc func_doc
         vector[shared_ptr[CDataType]] input_types
         shared_ptr[CDataType] output_type
+        CFunctionRegistry* registry
 
     CStatus RegisterScalarFunction(PyObject* function,
                                    function[CallbackUdf] wrapper, const CScalarUdfOptions& options)

python/pyarrow/src/arrow/python/udf.cc

@@ -115,8 +115,7 @@ Status RegisterScalarFunction(PyObject* user_function, ScalarUdfWrapperCallback
   kernel.mem_allocation = compute::MemAllocation::NO_PREALLOCATE;
   kernel.null_handling = compute::NullHandling::COMPUTED_NO_PREALLOCATE;
   RETURN_NOT_OK(scalar_func->AddKernel(std::move(kernel)));
-  auto registry = compute::GetFunctionRegistry();
-  RETURN_NOT_OK(registry->AddFunction(std::move(scalar_func)));
+  RETURN_NOT_OK(options.registry->AddFunction(std::move(scalar_func)));
   return Status::OK();
 }
 

python/pyarrow/src/arrow/python/udf.h

@@ -39,6 +39,7 @@ struct ARROW_PYTHON_EXPORT ScalarUdfOptions {
   compute::FunctionDoc func_doc;
   std::vector<std::shared_ptr<DataType>> input_types;
   std::shared_ptr<DataType> output_type;
+  compute::FunctionRegistry* registry;
 };
 
 /// \brief A context passed as the first argument of scalar UDF functions.

python/pyarrow/tests/test_udf.py

@@ -17,6 +17,7 @@
 
 
 import pytest
+import multiprocessing as mp
 
 import pyarrow as pa
 from pyarrow import compute as pc
@@ -504,3 +505,110 @@ def test_input_lifetime(unary_func_fixture):
     # Calling a UDF should not have kept `v` alive longer than required
     v = None
     assert proxy_pool.bytes_allocated() == 0
+
+
+def test_nested_function_registry():
+    def f1(ctx, x):
+        return pc.call_function("add", [x, 1],
+                                memory_pool=ctx.memory_pool)
+    func_name = "f1"
+    unary_doc = {"summary": "add function",
+                 "description": "test add function"}
+
+    default_registry = pc.function_registry()
+
+    registry1 = pc.FunctionRegistry(default_registry)
+
+    registry2 = pc.FunctionRegistry(registry1)
+
+    pc.register_scalar_function(f1,
+                                func_name,
+                                unary_doc,
+                                {"array": pa.int64()},
+                                pa.int64(),
+                                registry2)
+
+    assert registry2.get_function(func_name).name == func_name
+
+    error_msg = "No function registered with name: f1"
+    with pytest.raises(pa.lib.ArrowKeyError, match=error_msg):
+        registry1.get_function(func_name)
+
+    pc.register_scalar_function(f1,
+                                func_name,
+                                unary_doc,
+                                {"array": pa.int64()},
+                                pa.int64(),
+                                registry1)
+    assert registry1.get_function(func_name).name == func_name
+
+    pc.register_scalar_function(f1,
+                                func_name,
+                                unary_doc,
+                                {"array": pa.int64()},
+                                pa.int64(),
+                                default_registry)
+
+    assert default_registry.get_function(func_name).name == func_name
+
+
+def parallel_task1(data, q):
+    def f1(ctx, x):
+        return pc.call_function("add", [x, 10],
+                                memory_pool=ctx.memory_pool)
+    func_name = "f1"
+    unary_doc = {"summary": "add function",
+                 "description": "test add function"}
+
+    default_registry = pc.function_registry()
+
+    registry1 = pc.FunctionRegistry(default_registry)
+    pc.register_scalar_function(f1,
+                                func_name,
+                                unary_doc,
+                                {"array": pa.int64()},
+                                pa.int64(),
+                                registry1)
+    func = registry1.get_function(func_name)
+    result = func.call(data)
+    q.put(result)
+
+
+def parallel_task2(data, q):
+    def f1(ctx, x):
+        return pc.call_function("multiply", [x, 10],
+                                memory_pool=ctx.memory_pool)
+    func_name = "f1"
+    unary_doc = {"summary": "multiply function",
+                 "description": "test multiply function"}
+
+    default_registry = pc.function_registry()
+
+    registry1 = pc.FunctionRegistry(default_registry)
+    pc.register_scalar_function(f1,
+                                func_name,
+                                unary_doc,
+                                {"array": pa.int64()},
+                                pa.int64(),
+                                registry1)
+    func = registry1.get_function(func_name)
+    q.put(func.call(data))
+
+
+def test_udf_usage_by_scope():
+    # With support to custom function registration on global and
+    # nested function registries, user has the ability to make an
+    # scope for a registry and get some tasks done in a particular
+    # process and drop the registry once the process terminates.
+    ctx = mp.get_context('spawn')
+    q = ctx.Queue()
+    p1 = ctx.Process(target=parallel_task1, args=(
+        [pa.array([10, 20, 30], pa.int64())], q))
+    p1.start()
+    result = q.get()
+    p1.join()
+    p2 = ctx.Process(target=parallel_task2, args=([result], q))
+    p2.start()
+    final_result = q.get()
+    p2.join()
+    assert final_result == pa.array([200, 300, 400], pa.int64())