diff --git a/python/pyarrow/_compute.pyx b/python/pyarrow/_compute.pyx index 1c98bbfdea2c..c17870f3225a 100644 --- a/python/pyarrow/_compute.pyx +++ b/python/pyarrow/_compute.pyx @@ -3482,6 +3482,8 @@ def _register_user_defined_function(register_func, func, function_name, function if func_registry is None: c_func_registry = NULL + elif not isinstance(func_registry, FunctionRegistry): + raise TypeError("func_registry must be a FunctionRegistry") else: c_func_registry = (func_registry).registry @@ -3516,6 +3518,8 @@ def call_tabular_function(function_name, args=None, func_registry=None): c_func_name = tobytes(function_name) if func_registry is None: c_func_registry = NULL + elif not isinstance(func_registry, FunctionRegistry): + raise TypeError("func_registry must be a FunctionRegistry") else: c_func_registry = (func_registry).registry if args is None: diff --git a/python/pyarrow/tests/test_compute.py b/python/pyarrow/tests/test_compute.py index 8b2ad2b333fc..00e73cee9ddb 100644 --- a/python/pyarrow/tests/test_compute.py +++ b/python/pyarrow/tests/test_compute.py @@ -246,6 +246,31 @@ def test_list_functions(): assert "add" in pc.list_functions() +def test_call_tabular_function_rejects_invalid_registry(): + with pytest.raises(TypeError, + match="func_registry must be a FunctionRegistry"): + pc.call_tabular_function("", None, 1) + + +@pytest.mark.parametrize("register_function", [ + pc.register_scalar_function, + pc.register_vector_function, + pc.register_aggregate_function, + pc.register_tabular_function, +]) +def test_register_function_rejects_invalid_registry(register_function): + with pytest.raises(TypeError, + match="func_registry must be a FunctionRegistry"): + register_function( + func=lambda context: None, + function_name="invalid_registry", + function_doc={"summary": "", "description": ""}, + in_types={}, + out_type=pa.struct([]), + func_registry=1, + ) + + def _check_get_function(name, expected_func_cls, expected_ker_cls, min_num_kernels=1): func = pc.get_function(name)