westonpace commented on code in PR #14682:
URL: https://github.com/apache/arrow/pull/14682#discussion_r1057284406
##########
python/pyarrow/_compute.pyx:
##########
@@ -2568,31 +2580,53 @@ cdef inline CFunctionDoc _make_function_doc(dict
func_doc) except *:
return f_doc
-cdef object box_scalar_udf_context(const CScalarUdfContext& c_context):
- cdef ScalarUdfContext context = ScalarUdfContext.__new__(ScalarUdfContext)
+cdef object box_udf_context(const CUdfContext& c_context):
+ cdef UdfContext context = UdfContext.__new__(UdfContext)
context.init(c_context)
return context
-cdef _scalar_udf_callback(user_function, const CScalarUdfContext& c_context,
inputs):
+cdef _udf_callback(user_function, const CUdfContext& c_context, inputs):
"""
- Helper callback function used to wrap the ScalarUdfContext from Python to
C++
+ Helper callback function used to wrap the UdfContext from Python to C++
execution.
"""
- context = box_scalar_udf_context(c_context)
+ context = box_udf_context(c_context)
return user_function(context, *inputs)
-def _get_scalar_udf_context(memory_pool, batch_length):
- cdef CScalarUdfContext c_context
+def _get_udf_context(memory_pool, batch_length):
+ cdef CUdfContext c_context
c_context.pool = maybe_unbox_memory_pool(memory_pool)
c_context.batch_length = batch_length
- context = box_scalar_udf_context(c_context)
+ context = box_udf_context(c_context)
return context
-def register_scalar_function(func, function_name, function_doc, in_types,
- out_type):
+ctypedef CStatus (*CRegisterUdf)(PyObject* function, function[CallbackUdf]
wrapper,
+ const CUdfOptions& options,
CFunctionRegistry* registry)
+
+cdef class RegisterUdf(_Weakrefable):
+ cdef CRegisterUdf register_func
+
+ cdef void init(self, const CRegisterUdf register_func):
+ self.register_func = register_func
+
+
+cdef get_register_scalar_function():
+ cdef RegisterUdf reg = RegisterUdf.__new__(RegisterUdf)
+ reg.register_func = RegisterScalarFunction
+ return reg
+
+
+cdef get_register_tabular_function():
+ cdef RegisterUdf reg = RegisterUdf.__new__(RegisterUdf)
+ reg.register_func = RegisterTabularFunction
+ return reg
+
Review Comment:
My preference would be to keep this and, if we have something better when we
add aggregate udfs, we can switch to that then. That being said, I'm not 100%
certain I follow the arguments here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
