ARROW-12744: [C++][Compute] Add rounding kernel

cpp/src/arrow/compute/api_scalar.cc

@@ -53,6 +53,7 @@ struct EnumTraits<compute::JoinOptions::NullHandlingBehavior>
     return "<INVALID>";
   }
 };
+
 template <>
 struct EnumTraits<TimeUnit::type>
     : BasicEnumTraits<TimeUnit::type, TimeUnit::type::SECOND, TimeUnit::type::MILLI,
@@ -72,6 +73,7 @@ struct EnumTraits<TimeUnit::type>
     return "<INVALID>";
   }
 };
+
 template <>
 struct EnumTraits<compute::CompareOperator>
     : BasicEnumTraits<
@@ -136,6 +138,42 @@ struct EnumTraits<compute::AssumeTimezoneOptions::Nonexistent>
     return "<INVALID>";
   }
 };
+
+template <>
+struct EnumTraits<compute::RoundMode>
+    : BasicEnumTraits<compute::RoundMode, compute::RoundMode::DOWN,
+                      compute::RoundMode::UP, compute::RoundMode::TOWARDS_ZERO,
+                      compute::RoundMode::TOWARDS_INFINITY, compute::RoundMode::HALF_DOWN,
+                      compute::RoundMode::HALF_UP, compute::RoundMode::HALF_TOWARDS_ZERO,
+                      compute::RoundMode::HALF_TOWARDS_INFINITY,
+                      compute::RoundMode::HALF_TO_EVEN, compute::RoundMode::HALF_TO_ODD> {
+  static std::string name() { return "compute::RoundMode"; }
+  static std::string value_name(compute::RoundMode value) {
+    switch (value) {
+      case compute::RoundMode::DOWN:
+        return "DOWN";
+      case compute::RoundMode::UP:
+        return "UP";
+      case compute::RoundMode::TOWARDS_ZERO:
+        return "TOWARDS_ZERO";
+      case compute::RoundMode::TOWARDS_INFINITY:
+        return "TOWARDS_INFINITY";
+      case compute::RoundMode::HALF_DOWN:
+        return "HALF_DOWN";
+      case compute::RoundMode::HALF_UP:
+        return "HALF_UP";
+      case compute::RoundMode::HALF_TOWARDS_ZERO:
+        return "HALF_TOWARDS_ZERO";
+      case compute::RoundMode::HALF_TOWARDS_INFINITY:
+        return "HALF_TOWARDS_INFINITY";
+      case compute::RoundMode::HALF_TO_EVEN:
+        return "HALF_TO_EVEN";
+      case compute::RoundMode::HALF_TO_ODD:
+        return "HALF_TO_ODD";
+    }
+    return "<INVALID>";
+  }
+};
 }  // namespace internal
 
 namespace compute {
@@ -153,6 +191,12 @@ static auto kArithmeticOptionsType = GetFunctionOptionsType<ArithmeticOptions>(
 static auto kElementWiseAggregateOptionsType =
     GetFunctionOptionsType<ElementWiseAggregateOptions>(
         DataMember("skip_nulls", &ElementWiseAggregateOptions::skip_nulls));
+static auto kRoundOptionsType = GetFunctionOptionsType<RoundOptions>(
+    DataMember("ndigits", &RoundOptions::ndigits),
+    DataMember("round_mode", &RoundOptions::round_mode));
+static auto kRoundToMultipleOptionsType = GetFunctionOptionsType<RoundToMultipleOptions>(
+    DataMember("multiple", &RoundToMultipleOptions::multiple),
+    DataMember("round_mode", &RoundToMultipleOptions::round_mode));
 static auto kJoinOptionsType = GetFunctionOptionsType<JoinOptions>(
     DataMember("null_handling", &JoinOptions::null_handling),
     DataMember("null_replacement", &JoinOptions::null_replacement));
@@ -217,6 +261,30 @@ ElementWiseAggregateOptions::ElementWiseAggregateOptions(bool skip_nulls)
       skip_nulls(skip_nulls) {}
 constexpr char ElementWiseAggregateOptions::kTypeName[];
 
+RoundOptions::RoundOptions(int64_t ndigits, RoundMode round_mode)
+    : FunctionOptions(internal::kRoundOptionsType),
+      ndigits(ndigits),
+      round_mode(round_mode) {
+  static_assert(RoundMode::HALF_DOWN > RoundMode::DOWN &&
+                    RoundMode::HALF_DOWN > RoundMode::UP &&
+                    RoundMode::HALF_DOWN > RoundMode::TOWARDS_ZERO &&
+                    RoundMode::HALF_DOWN > RoundMode::TOWARDS_INFINITY &&
+                    RoundMode::HALF_DOWN < RoundMode::HALF_UP &&
+                    RoundMode::HALF_DOWN < RoundMode::HALF_TOWARDS_ZERO &&
+                    RoundMode::HALF_DOWN < RoundMode::HALF_TOWARDS_INFINITY &&
+                    RoundMode::HALF_DOWN < RoundMode::HALF_TO_EVEN &&
+                    RoundMode::HALF_DOWN < RoundMode::HALF_TO_ODD,
+                "Invalid order of round modes. Modes prefixed with HALF need to be "
+                "enumerated last with HALF_DOWN being the first among them.");
+}
+constexpr char RoundOptions::kTypeName[];
+
+RoundToMultipleOptions::RoundToMultipleOptions(double multiple, RoundMode round_mode)
+    : FunctionOptions(internal::kRoundToMultipleOptionsType),
+      multiple(multiple),
+      round_mode(round_mode) {}
+constexpr char RoundToMultipleOptions::kTypeName[];
+
 JoinOptions::JoinOptions(NullHandlingBehavior null_handling, std::string null_replacement)
     : FunctionOptions(internal::kJoinOptionsType),
       null_handling(null_handling),
@@ -352,6 +420,8 @@ namespace internal {
 void RegisterScalarOptions(FunctionRegistry* registry) {
   DCHECK_OK(registry->AddFunctionOptionsType(kArithmeticOptionsType));
   DCHECK_OK(registry->AddFunctionOptionsType(kElementWiseAggregateOptionsType));
+  DCHECK_OK(registry->AddFunctionOptionsType(kRoundOptionsType));
+  DCHECK_OK(registry->AddFunctionOptionsType(kRoundToMultipleOptionsType));
   DCHECK_OK(registry->AddFunctionOptionsType(kJoinOptionsType));
   DCHECK_OK(registry->AddFunctionOptionsType(kMatchSubstringOptionsType));
   DCHECK_OK(registry->AddFunctionOptionsType(kSplitOptionsType));
@@ -405,6 +475,15 @@ SCALAR_ARITHMETIC_UNARY(Log10, "log10", "log10_checked")
 SCALAR_ARITHMETIC_UNARY(Log2, "log2", "log2_checked")
 SCALAR_ARITHMETIC_UNARY(Log1p, "log1p", "log1p_checked")
 
+Result<Datum> Round(const Datum& arg, RoundOptions options, ExecContext* ctx) {
+  return CallFunction("round", {arg}, &options, ctx);
+}
+
+Result<Datum> RoundToMultiple(const Datum& arg, RoundToMultipleOptions options,
+                              ExecContext* ctx) {
+  return CallFunction("round_to_multiple", {arg}, &options, ctx);
+}
+
 #define SCALAR_ARITHMETIC_BINARY(NAME, REGISTRY_NAME, REGISTRY_CHECKED_NAME)           \
   Result<Datum> NAME(const Datum& left, const Datum& right, ArithmeticOptions options, \
                      ExecContext* ctx) {                                               \

cpp/src/arrow/compute/api_scalar.h

@@ -50,10 +50,58 @@ class ARROW_EXPORT ElementWiseAggregateOptions : public FunctionOptions {
   explicit ElementWiseAggregateOptions(bool skip_nulls = true);
   constexpr static char const kTypeName[] = "ElementWiseAggregateOptions";
   static ElementWiseAggregateOptions Defaults() { return ElementWiseAggregateOptions{}; }
-
   bool skip_nulls;
 };
 
+/// Rounding and tie-breaking modes for round compute functions.
+/// Additional details and examples are provided in compute.rst.
+enum class RoundMode : int8_t {
+  /// Round to nearest integer less than or equal in magnitude (aka "floor")
+  DOWN,
+  /// Round to nearest integer greater than or equal in magnitude (aka "ceil")
+  UP,
+  /// Get the integral part without fractional digits (aka "trunc")
+  TOWARDS_ZERO,
+  /// Round negative values with DOWN rule and positive values with UP rule
+  TOWARDS_INFINITY,
+  /// Round ties with DOWN rule
+  HALF_DOWN,
+  /// Round ties with UP rule
+  HALF_UP,
+  /// Round ties with TOWARDS_ZERO rule
+  HALF_TOWARDS_ZERO,
+  /// Round ties with TOWARDS_INFINITY rule
+  HALF_TOWARDS_INFINITY,
+  /// Round ties to nearest even integer
+  HALF_TO_EVEN,
+  /// Round ties to nearest odd integer
+  HALF_TO_ODD,
+};
+
+class ARROW_EXPORT RoundOptions : public FunctionOptions {
+ public:
+  explicit RoundOptions(int64_t ndigits = 0,
+                        RoundMode round_mode = RoundMode::HALF_TO_EVEN);
+  constexpr static char const kTypeName[] = "RoundOptions";
+  static RoundOptions Defaults() { return RoundOptions(); }
+  /// Rounding precision (number of digits to round to)
+  int64_t ndigits;
+  /// Rounding and tie-breaking mode
+  RoundMode round_mode;
+};
+
+class ARROW_EXPORT RoundToMultipleOptions : public FunctionOptions {
+ public:
+  explicit RoundToMultipleOptions(double multiple = 1.0,
+                                  RoundMode round_mode = RoundMode::HALF_TO_EVEN);
+  constexpr static char const kTypeName[] = "RoundToMultipleOptions";
+  static RoundToMultipleOptions Defaults() { return RoundToMultipleOptions(); }
+  /// Rounding scale (multiple to round to)
+  double multiple;
+  /// Rounding and tie-breaking mode
+  RoundMode round_mode;
+};
+
 /// Options for var_args_join.
 class ARROW_EXPORT JoinOptions : public FunctionOptions {
  public:
@@ -559,8 +607,9 @@ Result<Datum> Logb(const Datum& arg, const Datum& base,
                    ExecContext* ctx = NULLPTR);
 
 /// \brief Round to the nearest integer less than or equal in magnitude to the
-/// argument. Array values can be of arbitrary length. If argument is null the
-/// result will be null.
+/// argument.
+///
+/// If argument is null the result will be null.
 ///
 /// \param[in] arg the value to round
 /// \param[in] ctx the function execution context, optional
@@ -569,17 +618,19 @@ ARROW_EXPORT
 Result<Datum> Floor(const Datum& arg, ExecContext* ctx = NULLPTR);
 
 /// \brief Round to the nearest integer greater than or equal in magnitude to the
-/// argument. Array values can be of arbitrary length. If argument is null the
-/// result will be null.
+/// argument.
+///
+/// If argument is null the result will be null.
 ///
 /// \param[in] arg the value to round
 /// \param[in] ctx the function execution context, optional
 /// \return the rounded value
 ARROW_EXPORT
 Result<Datum> Ceil(const Datum& arg, ExecContext* ctx = NULLPTR);
 
-/// \brief Get the integral part without fractional digits. Array values can be
-/// of arbitrary length. If argument is null the result will be null.
+/// \brief Get the integral part without fractional digits.
+///
+/// If argument is null the result will be null.
 ///
 /// \param[in] arg the value to truncate
 /// \param[in] ctx the function execution context, optional
@@ -618,10 +669,35 @@ Result<Datum> MinElementWise(
 ///
 /// \param[in] arg the value to extract sign from
 /// \param[in] ctx the function execution context, optional
-/// \return the elementwise sign function
+/// \return the element-wise sign function
 ARROW_EXPORT
 Result<Datum> Sign(const Datum& arg, ExecContext* ctx = NULLPTR);
 
+/// \brief Round a value to a given precision.
+///
+/// If argument is null the result will be null.
+///
+/// \param[in] arg the value rounded
+/// \param[in] options rounding options (rounding mode and number of digits), optional
+/// \param[in] ctx the function execution context, optional
+/// \return the element-wise rounded value
+ARROW_EXPORT
+Result<Datum> Round(const Datum& arg, RoundOptions options = RoundOptions::Defaults(),
+                    ExecContext* ctx = NULLPTR);
+
+/// \brief Round a value to a given multiple.
+///
+/// If argument is null the result will be null.
+///
+/// \param[in] arg the value to round
+/// \param[in] options rounding options (rounding mode and multiple), optional
+/// \param[in] ctx the function execution context, optional
+/// \return the element-wise rounded value
+ARROW_EXPORT
+Result<Datum> RoundToMultiple(
+    const Datum& arg, RoundToMultipleOptions options = RoundToMultipleOptions::Defaults(),
+    ExecContext* ctx = NULLPTR);
+
 /// \brief Compare a numeric array with a scalar.
 ///
 /// \param[in] left datum to compare, must be an Array

cpp/src/arrow/compute/function.h

@@ -227,7 +227,7 @@ class ARROW_EXPORT Function {
   virtual Result<Datum> Execute(const std::vector<Datum>& args,
                                 const FunctionOptions* options, ExecContext* ctx) const;
 
-  /// \brief Returns a the default options for this function.
+  /// \brief Returns the default options for this function.
   ///
   /// Whatever option semantics a Function has, implementations must guarantee
   /// that default_options() is valid to pass to Execute as options.

cpp/src/arrow/compute/function_test.cc

@@ -58,6 +58,12 @@ TEST(FunctionOptions, Equality) {
   options.emplace_back(new IndexOptions(ScalarFromJSON(boolean(), "null")));
   options.emplace_back(new ArithmeticOptions());
   options.emplace_back(new ArithmeticOptions(/*check_overflow=*/true));
+  options.emplace_back(new RoundOptions());
+  options.emplace_back(
+      new RoundOptions(/*ndigits=*/2, /*round_mode=*/RoundMode::TOWARDS_INFINITY));
+  options.emplace_back(new RoundToMultipleOptions());
+  options.emplace_back(new RoundToMultipleOptions(
+      /*multiple=*/100, /*round_mode=*/RoundMode::TOWARDS_INFINITY));
   options.emplace_back(new ElementWiseAggregateOptions());
   options.emplace_back(new ElementWiseAggregateOptions(/*skip_nulls=*/false));
   options.emplace_back(new JoinOptions());

cpp/src/arrow/compute/kernels/codegen_internal.h

@@ -94,8 +94,8 @@ struct OptionsWrapper : public KernelState {
 /// KernelContext and the FunctionOptions as argument
 template <typename StateType, typename OptionsType>
 struct KernelStateFromFunctionOptions : public KernelState {
-  explicit KernelStateFromFunctionOptions(KernelContext* ctx, OptionsType state)
-      : state(StateType(ctx, std::move(state))) {}
+  explicit KernelStateFromFunctionOptions(KernelContext* ctx, OptionsType options)
+      : state(StateType(ctx, std::move(options))) {}
 
   static Result<std::unique_ptr<KernelState>> Init(KernelContext* ctx,
                                                    const KernelInitArgs& args) {