ARROW-15639 [C++][Python] UDF Scalar Function Implementation

cpp/examples/arrow/CMakeLists.txt

@@ -133,4 +133,7 @@ if(ARROW_PARQUET AND ARROW_DATASET)
 
   add_arrow_example(join_example EXTRA_LINK_LIBS ${DATASET_EXAMPLES_LINK_LIBS})
   add_dependencies(join-example parquet)
+
+  add_arrow_example(udf_example)
+
 endif()

cpp/examples/arrow/compute_register_example.cc

@@ -126,7 +126,7 @@ const cp::FunctionDoc func_doc{
 
 int main(int argc, char** argv) {
   const std::string name = "compute_register_example";
-  auto func = std::make_shared<cp::ScalarFunction>(name, cp::Arity::Unary(), &func_doc);
+  auto func = std::make_shared<cp::ScalarFunction>(name, cp::Arity::Unary(), func_doc);
   cp::ScalarKernel kernel({cp::InputType::Array(arrow::int64())}, arrow::int64(),
                           ExampleFunctionImpl);
   kernel.mem_allocation = cp::MemAllocation::NO_PREALLOCATE;

cpp/examples/arrow/udf_example.cc

@@ -0,0 +1,103 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+
+#include <cstdlib>
+#include <iostream>
+#include <memory>
+#include <string>
+#include <type_traits>
+#include <utility>
+#include <vector>
+
+// Demonstrate registering a user-defined Arrow compute function outside of the Arrow
+// source tree
+
+namespace cp = ::arrow::compute;
+
+#define ABORT_ON_FAILURE(expr)                     \
+  do {                                             \
+    arrow::Status status_ = (expr);                \
+    if (!status_.ok()) {                           \
+      std::cerr << status_.message() << std::endl; \
+      abort();                                     \
+    }                                              \
+  } while (0);
+
+template <typename TYPE,
+          typename = typename std::enable_if<arrow::is_number_type<TYPE>::value |
+                                             arrow::is_boolean_type<TYPE>::value |
+                                             arrow::is_temporal_type<TYPE>::value>::type>
+arrow::Result<std::shared_ptr<arrow::Array>> GetArrayDataSample(
+    const std::vector<typename TYPE::c_type>& values) {
+  using ArrowBuilderType = typename arrow::TypeTraits<TYPE>::BuilderType;
+  ArrowBuilderType builder;
+  ARROW_RETURN_NOT_OK(builder.Reserve(values.size()));
+  ARROW_RETURN_NOT_OK(builder.AppendValues(values));
+  return builder.Finish();
+}
+
+const cp::FunctionDoc func_doc{
+    "User-defined-function usage to demonstrate registering an out-of-tree function",
+    "returns x + y + z",
+    {"x", "y", "z"},
+    "UDFOptions"};
+
+arrow::Status SampleFunction(cp::KernelContext* ctx, const cp::ExecBatch& batch,
+                             arrow::Datum* out) {
+  // temp = x + y; return temp + z
+  ARROW_ASSIGN_OR_RAISE(auto temp, cp::CallFunction("add", {batch[0], batch[1]}));
+  return cp::CallFunction("add", {temp, batch[2]}).Value(out);
+}
+
+arrow::Status Execute() {
+  const std::string name = "add_three";
+  auto func = std::make_shared<cp::ScalarFunction>(name, cp::Arity::Ternary(), func_doc);
+  cp::ScalarKernel kernel(
+      {cp::InputType::Array(arrow::int64()), cp::InputType::Array(arrow::int64()),
+       cp::InputType::Array(arrow::int64())},
+      arrow::int64(), SampleFunction);
+
+  kernel.mem_allocation = cp::MemAllocation::NO_PREALLOCATE;
+  kernel.null_handling = cp::NullHandling::COMPUTED_NO_PREALLOCATE;
+
+  ARROW_RETURN_NOT_OK(func->AddKernel(std::move(kernel)));
+
+  auto registry = cp::GetFunctionRegistry();
+  ARROW_RETURN_NOT_OK(registry->AddFunction(std::move(func)));
+
+  ARROW_ASSIGN_OR_RAISE(auto x, GetArrayDataSample<arrow::Int64Type>({1, 2, 3}));
+  ARROW_ASSIGN_OR_RAISE(auto y, GetArrayDataSample<arrow::Int64Type>({4, 5, 6}));
+  ARROW_ASSIGN_OR_RAISE(auto z, GetArrayDataSample<arrow::Int64Type>({7, 8, 9}));
+
+  ARROW_ASSIGN_OR_RAISE(auto res, cp::CallFunction(name, {x, y, z}));
+  auto res_array = res.make_array();
+  std::cout << "Result" << std::endl;
+  std::cout << res_array->ToString() << std::endl;
+  return arrow::Status::OK();
+}
+
+int main(int argc, char** argv) {
+  auto status = Execute();
+  if (!status.ok()) {
+    std::cerr << "Error occurred : " << status.message() << std::endl;
+    return EXIT_FAILURE;
+  }
+  return EXIT_SUCCESS;
+}

cpp/src/arrow/compute/cast.cc

@@ -95,7 +95,7 @@ const FunctionDoc cast_doc{"Cast values to another data type",
 // to the standard SQL CAST(expr AS target_type)
 class CastMetaFunction : public MetaFunction {
  public:
-  CastMetaFunction() : MetaFunction("cast", Arity::Unary(), &cast_doc) {}
+  CastMetaFunction() : MetaFunction("cast", Arity::Unary(), cast_doc) {}
 
   Result<const CastOptions*> ValidateOptions(const FunctionOptions* options) const {
     auto cast_options = static_cast<const CastOptions*>(options);
@@ -153,7 +153,7 @@ CastOptions::CastOptions(bool safe)
 constexpr char CastOptions::kTypeName[];
 
 CastFunction::CastFunction(std::string name, Type::type out_type_id)
-    : ScalarFunction(std::move(name), Arity::Unary(), /*doc=*/nullptr),
+    : ScalarFunction(std::move(name), Arity::Unary(), FunctionDoc::Empty()),
       out_type_id_(out_type_id) {}
 
 Status CastFunction::AddKernel(Type::type in_type_id, ScalarKernel kernel) {

cpp/src/arrow/compute/exec_test.cc

@@ -681,8 +681,8 @@ class TestCallScalarFunction : public TestComputeInternals {
 
     // This function simply copies memory from the input argument into the
     // (preallocated) output
-    auto func =
-        std::make_shared<ScalarFunction>("test_copy", Arity::Unary(), /*doc=*/nullptr);
+    auto func = std::make_shared<ScalarFunction>("test_copy", Arity::Unary(),
+                                                 /*doc=*/FunctionDoc::Empty());
 
     // Add a few kernels. Our implementation only accepts arrays
     ASSERT_OK(func->AddKernel({InputType::Array(uint8())}, uint8(), ExecCopy));
@@ -691,8 +691,8 @@ class TestCallScalarFunction : public TestComputeInternals {
     ASSERT_OK(registry->AddFunction(func));
 
     // A version which doesn't want the executor to call PropagateNulls
-    auto func2 = std::make_shared<ScalarFunction>("test_copy_computed_bitmap",
-                                                  Arity::Unary(), /*doc=*/nullptr);
+    auto func2 = std::make_shared<ScalarFunction>(
+        "test_copy_computed_bitmap", Arity::Unary(), /*doc=*/FunctionDoc::Empty());
     ScalarKernel kernel({InputType::Array(uint8())}, uint8(), ExecComputedBitmap);
     kernel.null_handling = NullHandling::COMPUTED_PREALLOCATE;
     ASSERT_OK(func2->AddKernel(kernel));
@@ -705,9 +705,9 @@ class TestCallScalarFunction : public TestComputeInternals {
     // A function that allocates its own output memory. We have cases for both
     // non-preallocated data and non-preallocated validity bitmap
     auto f1 = std::make_shared<ScalarFunction>("test_nopre_data", Arity::Unary(),
-                                               /*doc=*/nullptr);
-    auto f2 = std::make_shared<ScalarFunction>("test_nopre_validity_or_data",
-                                               Arity::Unary(), /*doc=*/nullptr);
+                                               /*doc=*/FunctionDoc::Empty());
+    auto f2 = std::make_shared<ScalarFunction>(
+        "test_nopre_validity_or_data", Arity::Unary(), /*doc=*/FunctionDoc::Empty());
 
     ScalarKernel kernel({InputType::Array(uint8())}, uint8(), ExecNoPreallocatedData);
     kernel.mem_allocation = MemAllocation::NO_PREALLOCATE;
@@ -727,7 +727,7 @@ class TestCallScalarFunction : public TestComputeInternals {
     // This function's behavior depends on a static parameter that is made
     // available to the kernel's execution function through its Options object
     auto func = std::make_shared<ScalarFunction>("test_stateful", Arity::Unary(),
-                                                 /*doc=*/nullptr);
+                                                 /*doc=*/FunctionDoc::Empty());
 
     ScalarKernel kernel({InputType::Array(int32())}, int32(), ExecStateful, InitStateful);
     ASSERT_OK(func->AddKernel(kernel));
@@ -738,7 +738,7 @@ class TestCallScalarFunction : public TestComputeInternals {
     auto registry = GetFunctionRegistry();
 
     auto func = std::make_shared<ScalarFunction>("test_scalar_add_int32", Arity::Binary(),
-                                                 /*doc=*/nullptr);
+                                                 /*doc=*/FunctionDoc::Empty());
     ASSERT_OK(func->AddKernel({InputType::Scalar(int32()), InputType::Scalar(int32())},
                               int32(), ExecAddInt32));
     ASSERT_OK(registry->AddFunction(func));

cpp/src/arrow/compute/function.cc

@@ -290,9 +290,9 @@ Status ValidateFunctionDescription(const std::string& s) {
 }  // namespace
 
 Status Function::Validate() const {
-  if (!doc_->summary.empty()) {
+  if (!doc_.summary.empty()) {
     // Documentation given, check its contents
-    int arg_count = static_cast<int>(doc_->arg_names.size());
+    int arg_count = static_cast<int>(doc_.arg_names.size());
     // Some varargs functions allow 0 vararg, others expect at least 1,
     // hence the two possible values below.
     bool arg_count_match = (arg_count == arity_.num_args) ||
@@ -302,9 +302,9 @@ Status Function::Validate() const {
           "In function '", name_,
           "': ", "number of argument names for function documentation != function arity");
     }
-    Status st = ValidateFunctionSummary(doc_->summary);
+    Status st = ValidateFunctionSummary(doc_.summary);
     if (st.ok()) {
-      st &= ValidateFunctionDescription(doc_->description);
+      st &= ValidateFunctionDescription(doc_.description);
     }
     if (!st.ok()) {
       return st.WithMessage("In function '", name_, "': ", st.message());

cpp/src/arrow/compute/function.h

@@ -205,7 +205,7 @@ class ARROW_EXPORT Function {
   const Arity& arity() const { return arity_; }
 
   /// \brief Return the function documentation
-  const FunctionDoc& doc() const { return *doc_; }
+  const FunctionDoc& doc() const { return doc_; }
 
   /// \brief Returns the number of registered kernels for this function.
   virtual int num_kernels() const = 0;
@@ -244,12 +244,12 @@ class ARROW_EXPORT Function {
   virtual Status Validate() const;
 
  protected:
-  Function(std::string name, Function::Kind kind, const Arity& arity,
-           const FunctionDoc* doc, const FunctionOptions* default_options)
+  Function(std::string name, Function::Kind kind, const Arity& arity, FunctionDoc doc,
+           const FunctionOptions* default_options)
       : name_(std::move(name)),
         kind_(kind),
         arity_(arity),
-        doc_(doc ? doc : &FunctionDoc::Empty()),
+        doc_(std::move(doc)),
         default_options_(default_options) {}
 
   Status CheckArity(const std::vector<InputType>&) const;
@@ -258,7 +258,7 @@ class ARROW_EXPORT Function {
   std::string name_;
   Function::Kind kind_;
   Arity arity_;
-  const FunctionDoc* doc_;
+  const FunctionDoc doc_;
   const FunctionOptions* default_options_ = NULLPTR;
 };
 
@@ -279,9 +279,9 @@ class FunctionImpl : public Function {
   int num_kernels() const override { return static_cast<int>(kernels_.size()); }
 
  protected:
-  FunctionImpl(std::string name, Function::Kind kind, const Arity& arity,
-               const FunctionDoc* doc, const FunctionOptions* default_options)
-      : Function(std::move(name), kind, arity, doc, default_options) {}
+  FunctionImpl(std::string name, Function::Kind kind, const Arity& arity, FunctionDoc doc,
+               const FunctionOptions* default_options)
+      : Function(std::move(name), kind, arity, std::move(doc), default_options) {}
 
   std::vector<KernelType> kernels_;
 };
@@ -305,10 +305,10 @@ class ARROW_EXPORT ScalarFunction : public detail::FunctionImpl<ScalarKernel> {
  public:
   using KernelType = ScalarKernel;
 
-  ScalarFunction(std::string name, const Arity& arity, const FunctionDoc* doc,
+  ScalarFunction(std::string name, const Arity& arity, FunctionDoc doc,
                  const FunctionOptions* default_options = NULLPTR)
-      : detail::FunctionImpl<ScalarKernel>(std::move(name), Function::SCALAR, arity, doc,
-                                           default_options) {}
+      : detail::FunctionImpl<ScalarKernel>(std::move(name), Function::SCALAR, arity,
+                                           std::move(doc), default_options) {}
 
   /// \brief Add a kernel with given input/output types, no required state
   /// initialization, preallocation for fixed-width types, and default null
@@ -329,10 +329,10 @@ class ARROW_EXPORT VectorFunction : public detail::FunctionImpl<VectorKernel> {
  public:
   using KernelType = VectorKernel;
 
-  VectorFunction(std::string name, const Arity& arity, const FunctionDoc* doc,
+  VectorFunction(std::string name, const Arity& arity, FunctionDoc doc,
                  const FunctionOptions* default_options = NULLPTR)
-      : detail::FunctionImpl<VectorKernel>(std::move(name), Function::VECTOR, arity, doc,
-                                           default_options) {}
+      : detail::FunctionImpl<VectorKernel>(std::move(name), Function::VECTOR, arity,
+                                           std::move(doc), default_options) {}
 
   /// \brief Add a simple kernel with given input/output types, no required
   /// state initialization, no data preallocation, and no preallocation of the
@@ -350,10 +350,11 @@ class ARROW_EXPORT ScalarAggregateFunction
  public:
   using KernelType = ScalarAggregateKernel;
 
-  ScalarAggregateFunction(std::string name, const Arity& arity, const FunctionDoc* doc,
+  ScalarAggregateFunction(std::string name, const Arity& arity, FunctionDoc doc,
                           const FunctionOptions* default_options = NULLPTR)
-      : detail::FunctionImpl<ScalarAggregateKernel>(
-            std::move(name), Function::SCALAR_AGGREGATE, arity, doc, default_options) {}
+      : detail::FunctionImpl<ScalarAggregateKernel>(std::move(name),
+                                                    Function::SCALAR_AGGREGATE, arity,
+                                                    std::move(doc), default_options) {}
 
   /// \brief Add a kernel (function implementation). Returns error if the
   /// kernel's signature does not match the function's arity.
@@ -365,10 +366,11 @@ class ARROW_EXPORT HashAggregateFunction
  public:
   using KernelType = HashAggregateKernel;
 
-  HashAggregateFunction(std::string name, const Arity& arity, const FunctionDoc* doc,
+  HashAggregateFunction(std::string name, const Arity& arity, FunctionDoc doc,
                         const FunctionOptions* default_options = NULLPTR)
-      : detail::FunctionImpl<HashAggregateKernel>(
-            std::move(name), Function::HASH_AGGREGATE, arity, doc, default_options) {}
+      : detail::FunctionImpl<HashAggregateKernel>(std::move(name),
+                                                  Function::HASH_AGGREGATE, arity,
+                                                  std::move(doc), default_options) {}
 
   /// \brief Add a kernel (function implementation). Returns error if the
   /// kernel's signature does not match the function's arity.
@@ -392,9 +394,10 @@ class ARROW_EXPORT MetaFunction : public Function {
                                     const FunctionOptions* options,
                                     ExecContext* ctx) const = 0;
 
-  MetaFunction(std::string name, const Arity& arity, const FunctionDoc* doc,
+  MetaFunction(std::string name, const Arity& arity, FunctionDoc doc,
                const FunctionOptions* default_options = NULLPTR)
-      : Function(std::move(name), Function::META, arity, doc, default_options) {}
+      : Function(std::move(name), Function::META, arity, std::move(doc),
+                 default_options) {}
 };
 
 /// @}

cpp/src/arrow/compute/function_test.cc

@@ -179,8 +179,9 @@ TEST(Arity, Basics) {
 }
 
 TEST(ScalarFunction, Basics) {
-  ScalarFunction func("scalar_test", Arity::Binary(), /*doc=*/nullptr);
-  ScalarFunction varargs_func("varargs_test", Arity::VarArgs(1), /*doc=*/nullptr);
+  ScalarFunction func("scalar_test", Arity::Binary(), /*doc=*/FunctionDoc::Empty());
+  ScalarFunction varargs_func("varargs_test", Arity::VarArgs(1),
+                              /*doc=*/FunctionDoc::Empty());
 
   ASSERT_EQ("scalar_test", func.name());
   ASSERT_EQ(2, func.arity().num_args);
@@ -194,8 +195,9 @@ TEST(ScalarFunction, Basics) {
 }
 
 TEST(VectorFunction, Basics) {
-  VectorFunction func("vector_test", Arity::Binary(), /*doc=*/nullptr);
-  VectorFunction varargs_func("varargs_test", Arity::VarArgs(1), /*doc=*/nullptr);
+  VectorFunction func("vector_test", Arity::Binary(), /*doc=*/FunctionDoc::Empty());
+  VectorFunction varargs_func("varargs_test", Arity::VarArgs(1),
+                              /*doc=*/FunctionDoc::Empty());
 
   ASSERT_EQ("vector_test", func.name());
   ASSERT_EQ(2, func.arity().num_args);
@@ -260,15 +262,15 @@ void CheckAddDispatch(FunctionType* func) {
 }
 
 TEST(ScalarVectorFunction, DispatchExact) {
-  ScalarFunction func1("scalar_test", Arity::Binary(), /*doc=*/nullptr);
-  VectorFunction func2("vector_test", Arity::Binary(), /*doc=*/nullptr);
+  ScalarFunction func1("scalar_test", Arity::Binary(), /*doc=*/FunctionDoc::Empty());
+  VectorFunction func2("vector_test", Arity::Binary(), /*doc=*/FunctionDoc::Empty());
 
   CheckAddDispatch(&func1);
   CheckAddDispatch(&func2);
 }
 
 TEST(ArrayFunction, VarArgs) {
-  ScalarFunction va_func("va_test", Arity::VarArgs(1), /*doc=*/nullptr);
+  ScalarFunction va_func("va_test", Arity::VarArgs(1), /*doc=*/FunctionDoc::Empty());
 
   std::vector<InputType> va_args = {int8()};
 
@@ -294,7 +296,7 @@ TEST(ArrayFunction, VarArgs) {
 }
 
 TEST(ScalarAggregateFunction, Basics) {
-  ScalarAggregateFunction func("agg_test", Arity::Unary(), /*doc=*/nullptr);
+  ScalarAggregateFunction func("agg_test", Arity::Unary(), /*doc=*/FunctionDoc::Empty());
 
   ASSERT_EQ("agg_test", func.name());
   ASSERT_EQ(1, func.arity().num_args);
@@ -313,7 +315,7 @@ Status NoopMerge(KernelContext*, const KernelState&, KernelState*) {
 Status NoopFinalize(KernelContext*, Datum*) { return Status::OK(); }
 
 TEST(ScalarAggregateFunction, DispatchExact) {
-  ScalarAggregateFunction func("agg_test", Arity::Unary(), /*doc=*/nullptr);
+  ScalarAggregateFunction func("agg_test", Arity::Unary(), FunctionDoc::Empty());
 
   std::vector<InputType> in_args = {ValueDescr::Array(int8())};
   ScalarAggregateKernel kernel(std::move(in_args), int64(), NoopInit, NoopConsume,