# 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.

# cython: language_level = 3

import sys

from cpython.object cimport Py_LT, Py_EQ, Py_GT, Py_LE, Py_NE, Py_GE
from cython.operator cimport dereference as deref

from collections import namedtuple

from pyarrow.lib import frombytes, tobytes, ordered_dict, ArrowInvalid
from pyarrow.lib cimport *
from pyarrow.includes.common cimport *
from pyarrow.includes.libarrow cimport *
import pyarrow.lib as lib

from libcpp cimport bool as c_bool

import inspect
import numpy as np


cdef wrap_scalar_function(const shared_ptr[CFunction]& sp_func):
    """
    Wrap a C++ scalar Function in a ScalarFunction object.
    """
    cdef ScalarFunction func = ScalarFunction.__new__(ScalarFunction)
    func.init(sp_func)
    return func


cdef wrap_vector_function(const shared_ptr[CFunction]& sp_func):
    """
    Wrap a C++ vector Function in a VectorFunction object.
    """
    cdef VectorFunction func = VectorFunction.__new__(VectorFunction)
    func.init(sp_func)
    return func


cdef wrap_scalar_aggregate_function(const shared_ptr[CFunction]& sp_func):
    """
    Wrap a C++ aggregate Function in a ScalarAggregateFunction object.
    """
    cdef ScalarAggregateFunction func = \
        ScalarAggregateFunction.__new__(ScalarAggregateFunction)
    func.init(sp_func)
    return func


cdef wrap_hash_aggregate_function(const shared_ptr[CFunction]& sp_func):
    """
    Wrap a C++ aggregate Function in a HashAggregateFunction object.
    """
    cdef HashAggregateFunction func = \
        HashAggregateFunction.__new__(HashAggregateFunction)
    func.init(sp_func)
    return func


cdef wrap_meta_function(const shared_ptr[CFunction]& sp_func):
    """
    Wrap a C++ meta Function in a MetaFunction object.
    """
    cdef MetaFunction func = MetaFunction.__new__(MetaFunction)
    func.init(sp_func)
    return func


cdef wrap_function(const shared_ptr[CFunction]& sp_func):
    """
    Wrap a C++ Function in a Function object.

    This dispatches to specialized wrappers depending on the function kind.
    """
    if sp_func.get() == NULL:
        raise ValueError("Function was NULL")

    cdef FunctionKind c_kind = sp_func.get().kind()
    if c_kind == FunctionKind_SCALAR:
        return wrap_scalar_function(sp_func)
    elif c_kind == FunctionKind_VECTOR:
        return wrap_vector_function(sp_func)
    elif c_kind == FunctionKind_SCALAR_AGGREGATE:
        return wrap_scalar_aggregate_function(sp_func)
    elif c_kind == FunctionKind_HASH_AGGREGATE:
        return wrap_hash_aggregate_function(sp_func)
    elif c_kind == FunctionKind_META:
        return wrap_meta_function(sp_func)
    else:
        raise NotImplementedError("Unknown Function::Kind")


cdef wrap_scalar_kernel(const CScalarKernel* c_kernel):
    if c_kernel == NULL:
        raise ValueError("Kernel was NULL")
    cdef ScalarKernel kernel = ScalarKernel.__new__(ScalarKernel)
    kernel.init(c_kernel)
    return kernel


cdef wrap_vector_kernel(const CVectorKernel* c_kernel):
    if c_kernel == NULL:
        raise ValueError("Kernel was NULL")
    cdef VectorKernel kernel = VectorKernel.__new__(VectorKernel)
    kernel.init(c_kernel)
    return kernel


cdef wrap_scalar_aggregate_kernel(const CScalarAggregateKernel* c_kernel):
    if c_kernel == NULL:
        raise ValueError("Kernel was NULL")
    cdef ScalarAggregateKernel kernel = \
        ScalarAggregateKernel.__new__(ScalarAggregateKernel)
    kernel.init(c_kernel)
    return kernel


cdef wrap_hash_aggregate_kernel(const CHashAggregateKernel* c_kernel):
    if c_kernel == NULL:
        raise ValueError("Kernel was NULL")
    cdef HashAggregateKernel kernel = \
        HashAggregateKernel.__new__(HashAggregateKernel)
    kernel.init(c_kernel)
    return kernel


cdef class Kernel(_Weakrefable):
    """
    A kernel object.

    Kernels handle the execution of a Function for a certain signature.
    """

    def __init__(self):
        raise TypeError("Do not call {}'s constructor directly"
                        .format(self.__class__.__name__))


cdef class ScalarKernel(Kernel):
    cdef const CScalarKernel* kernel

    cdef void init(self, const CScalarKernel* kernel) except *:
        self.kernel = kernel

    def __repr__(self):
        return ("ScalarKernel<{}>"
                .format(frombytes(self.kernel.signature.get().ToString())))


cdef class VectorKernel(Kernel):
    cdef const CVectorKernel* kernel

    cdef void init(self, const CVectorKernel* kernel) except *:
        self.kernel = kernel

    def __repr__(self):
        return ("VectorKernel<{}>"
                .format(frombytes(self.kernel.signature.get().ToString())))


cdef class ScalarAggregateKernel(Kernel):
    cdef const CScalarAggregateKernel* kernel

    cdef void init(self, const CScalarAggregateKernel* kernel) except *:
        self.kernel = kernel

    def __repr__(self):
        return ("ScalarAggregateKernel<{}>"
                .format(frombytes(self.kernel.signature.get().ToString())))


cdef class HashAggregateKernel(Kernel):
    cdef const CHashAggregateKernel* kernel

    cdef void init(self, const CHashAggregateKernel* kernel) except *:
        self.kernel = kernel

    def __repr__(self):
        return ("HashAggregateKernel<{}>"
                .format(frombytes(self.kernel.signature.get().ToString())))


FunctionDoc = namedtuple(
    "FunctionDoc",
    ("summary", "description", "arg_names", "options_class",
     "options_required"))


cdef class Function(_Weakrefable):
    """
    A compute function.

    A function implements a certain logical computation over a range of
    possible input signatures.  Each signature accepts a range of input
    types and is implemented by a given Kernel.

    Functions can be of different kinds:

    * "scalar" functions apply an item-wise computation over all items
      of their inputs.  Each item in the output only depends on the values
      of the inputs at the same position.  Examples: addition, comparisons,
      string predicates...

    * "vector" functions apply a collection-wise computation, such that
      each item in the output may depend on the values of several items
      in each input.  Examples: dictionary encoding, sorting, extracting
      unique values...

    * "scalar_aggregate" functions reduce the dimensionality of the inputs by
      applying a reduction function.  Examples: sum, min_max, mode...

    * "hash_aggregate" functions apply a reduction function to an input
      subdivided by grouping criteria.  They may not be directly called.
      Examples: hash_sum, hash_min_max...

    * "meta" functions dispatch to other functions.
    """

    cdef:
        shared_ptr[CFunction] sp_func
        CFunction* base_func

    _kind_map = {
        FunctionKind_SCALAR: "scalar",
        FunctionKind_VECTOR: "vector",
        FunctionKind_SCALAR_AGGREGATE: "scalar_aggregate",
        FunctionKind_HASH_AGGREGATE: "hash_aggregate",
        FunctionKind_META: "meta",
    }

    def __init__(self):
        raise TypeError("Do not call {}'s constructor directly"
                        .format(self.__class__.__name__))

    cdef void init(self, const shared_ptr[CFunction]& sp_func) except *:
        self.sp_func = sp_func
        self.base_func = sp_func.get()

    def __repr__(self):
        return ("arrow.compute.Function<name={}, kind={}, "
                "arity={}, num_kernels={}>"
                .format(self.name, self.kind, self.arity, self.num_kernels))

    def __reduce__(self):
        # Reduction uses the global registry
        return get_function, (self.name,)

    @property
    def name(self):
        """
        The function name.
        """
        return frombytes(self.base_func.name())

    @property
    def arity(self):
        """
        The function arity.

        If Ellipsis (i.e. `...`) is returned, the function takes a variable
        number of arguments.
        """
        cdef CArity arity = self.base_func.arity()
        if arity.is_varargs:
            return ...
        else:
            return arity.num_args

    @property
    def kind(self):
        """
        The function kind.
        """
        cdef FunctionKind c_kind = self.base_func.kind()
        try:
            return self._kind_map[c_kind]
        except KeyError:
            raise NotImplementedError("Unknown Function::Kind")

    @property
    def _doc(self):
        """
        The C++-like function documentation (for internal use).
        """
        cdef CFunctionDoc c_doc = self.base_func.doc()
        return FunctionDoc(frombytes(c_doc.summary),
                           frombytes(c_doc.description),
                           [frombytes(s) for s in c_doc.arg_names],
                           frombytes(c_doc.options_class),
                           c_doc.options_required)

    @property
    def num_kernels(self):
        """
        The number of kernels implementing this function.
        """
        return self.base_func.num_kernels()

    def call(self, args, FunctionOptions options=None,
             MemoryPool memory_pool=None, length=None):
        """
        Call the function on the given arguments.

        Parameters
        ----------
        args : iterable
            The arguments to pass to the function.  Accepted types depend
            on the specific function.
        options : FunctionOptions, optional
            Options instance for executing this function.  This should have
            the right concrete options type.
        memory_pool : pyarrow.MemoryPool, optional
            If not passed, will allocate memory from the default memory pool.
        length : int, optional
            Batch size for execution, for nullary (no argument) functions. If
            not passed, will be inferred from passed data.
        """
        cdef:
            const CFunctionOptions* c_options = NULL
            CMemoryPool* pool = maybe_unbox_memory_pool(memory_pool)
            CExecContext c_exec_ctx = CExecContext(pool)
            CExecBatch c_batch
            CDatum result

        _pack_compute_args(args, &c_batch.values)

        if options is not None:
            c_options = options.get_options()

        if length is not None:
            c_batch.length = length
            with nogil:
                result = GetResultValue(
                    self.base_func.Execute(c_batch, c_options, &c_exec_ctx)
                )
        else:
            with nogil:
                result = GetResultValue(
                    self.base_func.Execute(c_batch.values, c_options,
                                           &c_exec_ctx)
                )

        return wrap_datum(result)


cdef class ScalarFunction(Function):
    cdef const CScalarFunction* func

    cdef void init(self, const shared_ptr[CFunction]& sp_func) except *:
        Function.init(self, sp_func)
        self.func = <const CScalarFunction*> sp_func.get()

    @property
    def kernels(self):
        """
        The kernels implementing this function.
        """
        cdef vector[const CScalarKernel*] kernels = self.func.kernels()
        return [wrap_scalar_kernel(k) for k in kernels]


cdef class VectorFunction(Function):
    cdef const CVectorFunction* func

    cdef void init(self, const shared_ptr[CFunction]& sp_func) except *:
        Function.init(self, sp_func)
        self.func = <const CVectorFunction*> sp_func.get()

    @property
    def kernels(self):
        """
        The kernels implementing this function.
        """
        cdef vector[const CVectorKernel*] kernels = self.func.kernels()
        return [wrap_vector_kernel(k) for k in kernels]


cdef class ScalarAggregateFunction(Function):
    cdef const CScalarAggregateFunction* func

    cdef void init(self, const shared_ptr[CFunction]& sp_func) except *:
        Function.init(self, sp_func)
        self.func = <const CScalarAggregateFunction*> sp_func.get()

    @property
    def kernels(self):
        """
        The kernels implementing this function.
        """
        cdef vector[const CScalarAggregateKernel*] kernels = \
            self.func.kernels()
        return [wrap_scalar_aggregate_kernel(k) for k in kernels]


cdef class HashAggregateFunction(Function):
    cdef const CHashAggregateFunction* func

    cdef void init(self, const shared_ptr[CFunction]& sp_func) except *:
        Function.init(self, sp_func)
        self.func = <const CHashAggregateFunction*> sp_func.get()

    @property
    def kernels(self):
        """
        The kernels implementing this function.
        """
        cdef vector[const CHashAggregateKernel*] kernels = self.func.kernels()
        return [wrap_hash_aggregate_kernel(k) for k in kernels]


cdef class MetaFunction(Function):
    cdef const CMetaFunction* func

    cdef void init(self, const shared_ptr[CFunction]& sp_func) except *:
        Function.init(self, sp_func)
        self.func = <const CMetaFunction*> sp_func.get()

    # Since num_kernels is exposed, also expose a kernels property
    @property
    def kernels(self):
        """
        The kernels implementing this function.
        """
        return []


cdef _pack_compute_args(object values, vector[CDatum]* out):
    for val in values:
        if isinstance(val, (list, np.ndarray)):
            val = lib.asarray(val)

        if isinstance(val, Array):
            out.push_back(CDatum((<Array> val).sp_array))
            continue
        elif isinstance(val, ChunkedArray):
            out.push_back(CDatum((<ChunkedArray> val).sp_chunked_array))
            continue
        elif isinstance(val, Scalar):
            out.push_back(CDatum((<Scalar> val).unwrap()))
            continue
        elif isinstance(val, RecordBatch):
            out.push_back(CDatum((<RecordBatch> val).sp_batch))
            continue
        elif isinstance(val, Table):
            out.push_back(CDatum((<Table> val).sp_table))
            continue
        else:
            # Is it a Python scalar?
            try:
                scal = lib.scalar(val)
            except Exception:
                # Raise dedicated error below
                pass
            else:
                out.push_back(CDatum((<Scalar> scal).unwrap()))
                continue

        raise TypeError(f"Got unexpected argument type {type(val)} "
                        "for compute function")


cdef class FunctionRegistry(_Weakrefable):
    cdef CFunctionRegistry* registry

    def __init__(self):
        self.registry = GetFunctionRegistry()

    def list_functions(self):
        """
        Return all function names in the registry.
        """
        cdef vector[c_string] names = self.registry.GetFunctionNames()
        return [frombytes(name) for name in names]

    def get_function(self, name):
        """
        Look up a function by name in the registry.

        Parameters
        ----------
        name : str
            The name of the function to lookup
        """
        cdef:
            c_string c_name = tobytes(name)
            shared_ptr[CFunction] func
        with nogil:
            func = GetResultValue(self.registry.GetFunction(c_name))
        return wrap_function(func)


cdef FunctionRegistry _global_func_registry = FunctionRegistry()


def function_registry():
    return _global_func_registry


def get_function(name):
    """
    Get a function by name.

    The function is looked up in the global registry
    (as returned by `function_registry()`).

    Parameters
    ----------
    name : str
        The name of the function to lookup
    """
    return _global_func_registry.get_function(name)


def list_functions():
    """
    Return all function names in the global registry.
    """
    return _global_func_registry.list_functions()


def call_function(name, args, options=None, memory_pool=None, length=None):
    """
    Call a named function.

    The function is looked up in the global registry
    (as returned by `function_registry()`).

    Parameters
    ----------
    name : str
        The name of the function to call.
    args : list
        The arguments to the function.
    options : optional
        options provided to the function.
    memory_pool : MemoryPool, optional
        memory pool to use for allocations during function execution.
    length : int, optional
        Batch size for execution, for nullary (no argument) functions. If not
        passed, inferred from data.
    """
    func = _global_func_registry.get_function(name)
    return func.call(args, options=options, memory_pool=memory_pool,
                     length=length)


cdef class FunctionOptions(_Weakrefable):
    __slots__ = ()  # avoid mistakingly creating attributes

    cdef const CFunctionOptions* get_options(self) except NULL:
        return self.wrapped.get()

    cdef void init(self, const shared_ptr[CFunctionOptions]& sp):
        self.wrapped = sp

    cdef inline shared_ptr[CFunctionOptions] unwrap(self):
        return self.wrapped

    def serialize(self):
        cdef:
            CResult[shared_ptr[CBuffer]] res = self.get_options().Serialize()
            shared_ptr[CBuffer] c_buf = GetResultValue(res)
        return pyarrow_wrap_buffer(c_buf)

    @staticmethod
    def deserialize(buf):
        """
        Deserialize options for a function.

        Parameters
        ----------
        buf : Buffer
            The buffer containing the data to deserialize.
        """
        cdef:
            shared_ptr[CBuffer] c_buf = pyarrow_unwrap_buffer(buf)
            CResult[unique_ptr[CFunctionOptions]] maybe_options = \
                DeserializeFunctionOptions(deref(c_buf))
            shared_ptr[CFunctionOptions] c_options
        c_options = to_shared(GetResultValue(move(maybe_options)))
        type_name = frombytes(c_options.get().options_type().type_name())
        module = globals()
        if type_name not in module:
            raise ValueError(f'Cannot deserialize "{type_name}"')
        klass = module[type_name]
        options = klass.__new__(klass)
        (<FunctionOptions> options).init(c_options)
        return options

    def __repr__(self):
        type_name = self.__class__.__name__
        # Remove {} so we can use our own braces
        string_repr = frombytes(self.get_options().ToString())[1:-1]
        return f"{type_name}({string_repr})"

    def __eq__(self, FunctionOptions other):
        return self.get_options().Equals(deref(other.get_options()))


def _raise_invalid_function_option(value, description, *,
                                   exception_class=ValueError):
    raise exception_class(f"\"{value}\" is not a valid {description}")


# NOTE:
# To properly expose the constructor signature of FunctionOptions
# subclasses, we use a two-level inheritance:
# 1. a C extension class that implements option validation and setting
#    (won't expose function signatures because of
#     https://github.com/cython/cython/issues/3873)
# 2. a Python derived class that implements the constructor

cdef class _CastOptions(FunctionOptions):
    cdef CCastOptions* options

    cdef void init(self, const shared_ptr[CFunctionOptions]& sp):
        FunctionOptions.init(self, sp)
        self.options = <CCastOptions*> self.wrapped.get()

    def _set_options(self, DataType target_type, allow_int_overflow,
                     allow_time_truncate, allow_time_overflow,
                     allow_decimal_truncate, allow_float_truncate,
                     allow_invalid_utf8):
        cdef:
            shared_ptr[CCastOptions] wrapped = make_shared[CCastOptions]()
        self.init(<shared_ptr[CFunctionOptions]> wrapped)
        self._set_type(target_type)
        if allow_int_overflow is not None:
            self.allow_int_overflow = allow_int_overflow
        if allow_time_truncate is not None:
            self.allow_time_truncate = allow_time_truncate
        if allow_time_overflow is not None:
            self.allow_time_overflow = allow_time_overflow
        if allow_decimal_truncate is not None:
            self.allow_decimal_truncate = allow_decimal_truncate
        if allow_float_truncate is not None:
            self.allow_float_truncate = allow_float_truncate
        if allow_invalid_utf8 is not None:
            self.allow_invalid_utf8 = allow_invalid_utf8

    def _set_type(self, target_type=None):
        if target_type is not None:
            deref(self.options).to_type = \
                (<DataType> ensure_type(target_type)).sp_type

    def _set_safe(self):
        self.init(shared_ptr[CFunctionOptions](
            new CCastOptions(CCastOptions.Safe())))

    def _set_unsafe(self):
        self.init(shared_ptr[CFunctionOptions](
            new CCastOptions(CCastOptions.Unsafe())))

    def is_safe(self):
        return not (deref(self.options).allow_int_overflow or
                    deref(self.options).allow_time_truncate or
                    deref(self.options).allow_time_overflow or
                    deref(self.options).allow_decimal_truncate or
                    deref(self.options).allow_float_truncate or
                    deref(self.options).allow_invalid_utf8)

    @property
    def allow_int_overflow(self):
        return deref(self.options).allow_int_overflow

    @allow_int_overflow.setter
    def allow_int_overflow(self, c_bool flag):
        deref(self.options).allow_int_overflow = flag

    @property
    def allow_time_truncate(self):
        return deref(self.options).allow_time_truncate

    @allow_time_truncate.setter
    def allow_time_truncate(self, c_bool flag):
        deref(self.options).allow_time_truncate = flag

    @property
    def allow_time_overflow(self):
        return deref(self.options).allow_time_overflow

    @allow_time_overflow.setter
    def allow_time_overflow(self, c_bool flag):
        deref(self.options).allow_time_overflow = flag

    @property
    def allow_decimal_truncate(self):
        return deref(self.options).allow_decimal_truncate

    @allow_decimal_truncate.setter
    def allow_decimal_truncate(self, c_bool flag):
        deref(self.options).allow_decimal_truncate = flag

    @property
    def allow_float_truncate(self):
        return deref(self.options).allow_float_truncate

    @allow_float_truncate.setter
    def allow_float_truncate(self, c_bool flag):
        deref(self.options).allow_float_truncate = flag

    @property
    def allow_invalid_utf8(self):
        return deref(self.options).allow_invalid_utf8

    @allow_invalid_utf8.setter
    def allow_invalid_utf8(self, c_bool flag):
        deref(self.options).allow_invalid_utf8 = flag


class CastOptions(_CastOptions):
    """
    Options for the `cast` function.

    Parameters
    ----------
    target_type : DataType, optional
        The PyArrow type to cast to.
    allow_int_overflow : bool, default False
        Whether integer overflow is allowed when casting.
    allow_time_truncate : bool, default False
        Whether time precision truncation is allowed when casting.
    allow_time_overflow : bool, default False
        Whether date/time range overflow is allowed when casting.
    allow_decimal_truncate : bool, default False
        Whether decimal precision truncation is allowed when casting.
    allow_float_truncate : bool, default False
        Whether floating-point precision truncation is allowed when casting.
    allow_invalid_utf8 : bool, default False
        Whether producing invalid utf8 data is allowed when casting.
    """

    def __init__(self, target_type=None, *, allow_int_overflow=None,
                 allow_time_truncate=None, allow_time_overflow=None,
                 allow_decimal_truncate=None, allow_float_truncate=None,
                 allow_invalid_utf8=None):
        self._set_options(target_type, allow_int_overflow, allow_time_truncate,
                          allow_time_overflow, allow_decimal_truncate,
                          allow_float_truncate, allow_invalid_utf8)

    @staticmethod
    def safe(target_type=None):
        """"
        Create a CastOptions for a safe cast.

        Parameters
        ----------
        target_type : optional
            Target cast type for the safe cast.
        """
        self = CastOptions()
        self._set_safe()
        self._set_type(target_type)
        return self

    @staticmethod
    def unsafe(target_type=None):
        """"
        Create a CastOptions for an unsafe cast.

        Parameters
        ----------
        target_type : optional
            Target cast type for the unsafe cast.
        """
        self = CastOptions()
        self._set_unsafe()
        self._set_type(target_type)
        return self


def _skip_nulls_doc():
    # (note the weird indent because of how the string is inserted
    #  by callers)
    return """skip_nulls : bool, default True
        Whether to skip (ignore) nulls in the input.
        If False, any null in the input forces the output to null.
"""


def _min_count_doc(*, default):
    return f"""min_count : int, default {default}
        Minimum number of non-null values in the input.  If the number
        of non-null values is below `min_count`, the output is null.
"""


cdef class _ElementWiseAggregateOptions(FunctionOptions):
    def _set_options(self, skip_nulls):
        self.wrapped.reset(new CElementWiseAggregateOptions(skip_nulls))


class ElementWiseAggregateOptions(_ElementWiseAggregateOptions):
    __doc__ = f"""
    Options for element-wise aggregate functions.

    Parameters
    ----------
    {_skip_nulls_doc()}
    """

    def __init__(self, *, skip_nulls=True):
        self._set_options(skip_nulls)


cdef CRoundMode unwrap_round_mode(round_mode) except *:
    if round_mode == "down":
        return CRoundMode_DOWN
    elif round_mode == "up":
        return CRoundMode_UP
    elif round_mode == "towards_zero":
        return CRoundMode_TOWARDS_ZERO
    elif round_mode == "towards_infinity":
        return CRoundMode_TOWARDS_INFINITY
    elif round_mode == "half_down":
        return CRoundMode_HALF_DOWN
    elif round_mode == "half_up":
        return CRoundMode_HALF_UP
    elif round_mode == "half_towards_zero":
        return CRoundMode_HALF_TOWARDS_ZERO
    elif round_mode == "half_towards_infinity":
        return CRoundMode_HALF_TOWARDS_INFINITY
    elif round_mode == "half_to_even":
        return CRoundMode_HALF_TO_EVEN
    elif round_mode == "half_to_odd":
        return CRoundMode_HALF_TO_ODD
    _raise_invalid_function_option(round_mode, "round mode")


cdef class _RoundOptions(FunctionOptions):
    def _set_options(self, ndigits, round_mode):
        self.wrapped.reset(
            new CRoundOptions(ndigits, unwrap_round_mode(round_mode))
        )


class RoundOptions(_RoundOptions):
    """
    Options for rounding numbers.

    Parameters
    ----------
    ndigits : int, default 0
        Number of fractional digits to round to.
    round_mode : str, default "half_to_even"
        Rounding and tie-breaking mode.
        Accepted values are "down", "up", "towards_zero", "towards_infinity",
        "half_down", "half_up", "half_towards_zero", "half_towards_infinity",
        "half_to_even", "half_to_odd".
    """

    def __init__(self, ndigits=0, round_mode="half_to_even"):
        self._set_options(ndigits, round_mode)


cdef CCalendarUnit unwrap_round_temporal_unit(unit) except *:
    if unit == "nanosecond":
        return CCalendarUnit_NANOSECOND
    elif unit == "microsecond":
        return CCalendarUnit_MICROSECOND
    elif unit == "millisecond":
        return CCalendarUnit_MILLISECOND
    elif unit == "second":
        return CCalendarUnit_SECOND
    elif unit == "minute":
        return CCalendarUnit_MINUTE
    elif unit == "hour":
        return CCalendarUnit_HOUR
    elif unit == "day":
        return CCalendarUnit_DAY
    elif unit == "week":
        return CCalendarUnit_WEEK
    elif unit == "month":
        return CCalendarUnit_MONTH
    elif unit == "quarter":
        return CCalendarUnit_QUARTER
    elif unit == "year":
        return CCalendarUnit_YEAR
    _raise_invalid_function_option(unit, "Calendar unit")


cdef class _RoundTemporalOptions(FunctionOptions):
    def _set_options(self, multiple, unit, week_starts_monday,
                     ceil_is_strictly_greater, calendar_based_origin):
        self.wrapped.reset(
            new CRoundTemporalOptions(
                multiple, unwrap_round_temporal_unit(unit),
                week_starts_monday, ceil_is_strictly_greater,
                calendar_based_origin)
        )


class RoundTemporalOptions(_RoundTemporalOptions):
    """
    Options for rounding temporal values.

    Parameters
    ----------
    multiple : int, default 1
        Number of units to round to.
    unit : str, default "day"
        The unit in which `multiple` is expressed.
        Accepted values are "year", "quarter", "month", "week", "day",
        "hour", "minute", "second", "millisecond", "microsecond",
        "nanosecond".
    week_starts_monday : bool, default True
        If True, weeks start on Monday; if False, on Sunday.
    ceil_is_strictly_greater : bool, default False
        If True, ceil returns a rounded value that is strictly greater than the
        input. For example: ceiling 1970-01-01T00:00:00 to 3 hours would
        yield 1970-01-01T03:00:00 if set to True and 1970-01-01T00:00:00
        if set to False.
        This applies to the ceil_temporal function only.
    calendar_based_origin : bool, default False
        By default, the origin is 1970-01-01T00:00:00. By setting this to True,
        rounding origin will be beginning of one less precise calendar unit.
        E.g.: rounding to hours will use beginning of day as origin.

        By default time is rounded to a multiple of units since
        1970-01-01T00:00:00. By setting calendar_based_origin to true,
        time will be rounded to number of units since the last greater
        calendar unit.
        For example: rounding to multiple of days since the beginning of the
        month or to hours since the beginning of the day.
        Exceptions: week and quarter are not used as greater units,
        therefore days will be rounded to the beginning of the month not
        week. Greater unit of week is a year.
        Note that ceiling and rounding might change sorting order of an array
        near greater unit change. For example rounding YYYY-mm-dd 23:00:00 to
        5 hours will ceil and round to YYYY-mm-dd+1 01:00:00 and floor to
        YYYY-mm-dd 20:00:00. On the other hand YYYY-mm-dd+1 00:00:00 will
        ceil, round and floor to YYYY-mm-dd+1 00:00:00. This can break the
        order of an already ordered array.

    """

    def __init__(self, multiple=1, unit="day", *, week_starts_monday=True,
                 ceil_is_strictly_greater=False,
                 calendar_based_origin=False):
        self._set_options(multiple, unit, week_starts_monday,
                          ceil_is_strictly_greater,
                          calendar_based_origin)


cdef class _RoundToMultipleOptions(FunctionOptions):
    def _set_options(self, multiple, round_mode):
        if not isinstance(multiple, Scalar):
            try:
                multiple = lib.scalar(multiple)
            except Exception:
                _raise_invalid_function_option(
                    multiple, "multiple type for RoundToMultipleOptions",
                    exception_class=TypeError)

        self.wrapped.reset(
            new CRoundToMultipleOptions(
                pyarrow_unwrap_scalar(multiple), unwrap_round_mode(round_mode))
        )


class RoundToMultipleOptions(_RoundToMultipleOptions):
    """
    Options for rounding numbers to a multiple.

    Parameters
    ----------
    multiple : numeric scalar, default 1.0
        Multiple to round to. Should be a scalar of a type compatible
        with the argument to be rounded.
    round_mode : str, default "half_to_even"
        Rounding and tie-breaking mode.
        Accepted values are "down", "up", "towards_zero", "towards_infinity",
        "half_down", "half_up", "half_towards_zero", "half_towards_infinity",
        "half_to_even", "half_to_odd".
    """

    def __init__(self, multiple=1.0, round_mode="half_to_even"):
        self._set_options(multiple, round_mode)


cdef class _JoinOptions(FunctionOptions):
    _null_handling_map = {
        "emit_null": CJoinNullHandlingBehavior_EMIT_NULL,
        "skip": CJoinNullHandlingBehavior_SKIP,
        "replace": CJoinNullHandlingBehavior_REPLACE,
    }

    def _set_options(self, null_handling, null_replacement):
        try:
            self.wrapped.reset(
                new CJoinOptions(self._null_handling_map[null_handling],
                                 tobytes(null_replacement))
            )
        except KeyError:
            _raise_invalid_function_option(null_handling, "null handling")


class JoinOptions(_JoinOptions):
    """
    Options for the `binary_join_element_wise` function.

    Parameters
    ----------
    null_handling : str, default "emit_null"
        How to handle null values in the inputs.
        Accepted values are "emit_null", "skip", "replace".
    null_replacement : str, default ""
        Replacement string to emit for null inputs if `null_handling`
        is "replace".
    """

    def __init__(self, null_handling="emit_null", null_replacement=""):
        self._set_options(null_handling, null_replacement)


cdef class _MatchSubstringOptions(FunctionOptions):
    def _set_options(self, pattern, ignore_case):
        self.wrapped.reset(
            new CMatchSubstringOptions(tobytes(pattern), ignore_case)
        )


class MatchSubstringOptions(_MatchSubstringOptions):
    """
    Options for looking for a substring.

    Parameters
    ----------
    pattern : str
        Substring pattern to look for inside input values.
    ignore_case : bool, default False
        Whether to perform a case-insensitive match.
    """

    def __init__(self, pattern, *, ignore_case=False):
        self._set_options(pattern, ignore_case)


cdef class _PadOptions(FunctionOptions):
    def _set_options(self, width, padding):
        self.wrapped.reset(new CPadOptions(width, tobytes(padding)))


class PadOptions(_PadOptions):
    """
    Options for padding strings.

    Parameters
    ----------
    width : int
        Desired string length.
    padding : str, default " "
        What to pad the string with. Should be one byte or codepoint.
    """

    def __init__(self, width, padding=' '):
        self._set_options(width, padding)


cdef class _TrimOptions(FunctionOptions):
    def _set_options(self, characters):
        self.wrapped.reset(new CTrimOptions(tobytes(characters)))


class TrimOptions(_TrimOptions):
    """
    Options for trimming characters from strings.

    Parameters
    ----------
    characters : str
        Individual characters to be trimmed from the string.
    """

    def __init__(self, characters):
        self._set_options(tobytes(characters))


cdef class _ReplaceSubstringOptions(FunctionOptions):
    def _set_options(self, pattern, replacement, max_replacements):
        self.wrapped.reset(
            new CReplaceSubstringOptions(tobytes(pattern),
                                         tobytes(replacement),
                                         max_replacements)
        )


class ReplaceSubstringOptions(_ReplaceSubstringOptions):
    """
    Options for replacing matched substrings.

    Parameters
    ----------
    pattern : str
        Substring pattern to look for inside input values.
    replacement : str
        What to replace the pattern with.
    max_replacements : int or None, default None
        The maximum number of strings to replace in each
        input value (unlimited if None).
    """

    def __init__(self, pattern, replacement, *, max_replacements=None):
        if max_replacements is None:
            max_replacements = -1
        self._set_options(pattern, replacement, max_replacements)


cdef class _ExtractRegexOptions(FunctionOptions):
    def _set_options(self, pattern):
        self.wrapped.reset(new CExtractRegexOptions(tobytes(pattern)))


class ExtractRegexOptions(_ExtractRegexOptions):
    """
    Options for the `extract_regex` function.

    Parameters
    ----------
    pattern : str
        Regular expression with named capture fields.
    """

    def __init__(self, pattern):
        self._set_options(pattern)


cdef class _SliceOptions(FunctionOptions):
    def _set_options(self, start, stop, step):
        self.wrapped.reset(new CSliceOptions(start, stop, step))


class SliceOptions(_SliceOptions):
    """
    Options for slicing.

    Parameters
    ----------
    start : int
        Index to start slicing at (inclusive).
    stop : int or None, default None
        If given, index to stop slicing at (exclusive).
        If not given, slicing will stop at the end.
    step : int, default 1
        Slice step.
    """

    def __init__(self, start, stop=None, step=1):
        if stop is None:
            stop = sys.maxsize
        self._set_options(start, stop, step)


cdef class _ListSliceOptions(FunctionOptions):
    cpdef _set_options(self, start, stop=None, step=1, return_fixed_size_list=None):
        cdef:
            CListSliceOptions* opts
        opts = new CListSliceOptions(
            start,
            <optional[int64_t]>nullopt if stop is None
            else <optional[int64_t]>(<int64_t>stop),
            step,
            <optional[c_bool]>nullopt if return_fixed_size_list is None
            else <optional[c_bool]>(<c_bool>return_fixed_size_list)
        )
        self.wrapped.reset(opts)


class ListSliceOptions(_ListSliceOptions):
    """
    Options for list array slicing.

    Parameters
    ----------
    start : int
        Index to start slicing inner list elements (inclusive).
    stop : Optional[int], default None
        If given, index to stop slicing at (exclusive).
        If not given, slicing will stop at the end. (NotImplemented)
    step : int, default 1
        Slice step.
    return_fixed_size_list : Optional[bool], default None
        Whether to return a FixedSizeListArray. If true _and_ stop is after
        a list element's length, nulls will be appended to create the
        requested slice size. The default of `None` will return the same
        type which was passed in.
    """

    def __init__(self, start, stop=None, step=1, return_fixed_size_list=None):
        self._set_options(start, stop, step, return_fixed_size_list)


cdef class _ReplaceSliceOptions(FunctionOptions):
    def _set_options(self, start, stop, replacement):
        self.wrapped.reset(
            new CReplaceSliceOptions(start, stop, tobytes(replacement))
        )


class ReplaceSliceOptions(_ReplaceSliceOptions):
    """
    Options for replacing slices.

    Parameters
    ----------
    start : int
        Index to start slicing at (inclusive).
    stop : int
        Index to stop slicing at (exclusive).
    replacement : str
        What to replace the slice with.
    """

    def __init__(self, start, stop, replacement):
        self._set_options(start, stop, replacement)


cdef class _FilterOptions(FunctionOptions):
    _null_selection_map = {
        "drop": CFilterNullSelectionBehavior_DROP,
        "emit_null": CFilterNullSelectionBehavior_EMIT_NULL,
    }

    def _set_options(self, null_selection_behavior):
        try:
            self.wrapped.reset(
                new CFilterOptions(
                    self._null_selection_map[null_selection_behavior]
                )
            )
        except KeyError:
            _raise_invalid_function_option(null_selection_behavior,
                                           "null selection behavior")


class FilterOptions(_FilterOptions):
    """
    Options for selecting with a boolean filter.

    Parameters
    ----------
    null_selection_behavior : str, default "drop"
        How to handle nulls in the selection filter.
        Accepted values are "drop", "emit_null".
    """

    def __init__(self, null_selection_behavior="drop"):
        self._set_options(null_selection_behavior)


cdef class _DictionaryEncodeOptions(FunctionOptions):
    _null_encoding_map = {
        "encode": CDictionaryEncodeNullEncodingBehavior_ENCODE,
        "mask": CDictionaryEncodeNullEncodingBehavior_MASK,
    }

    def _set_options(self, null_encoding):
        try:
            self.wrapped.reset(
                new CDictionaryEncodeOptions(
                    self._null_encoding_map[null_encoding]
                )
            )
        except KeyError:
            _raise_invalid_function_option(null_encoding, "null encoding")


class DictionaryEncodeOptions(_DictionaryEncodeOptions):
    """
    Options for dictionary encoding.

    Parameters
    ----------
    null_encoding : str, default "mask"
        How to encode nulls in the input.
        Accepted values are "mask" (null inputs emit a null in the indices
        array), "encode" (null inputs emit a non-null index pointing to
        a null value in the dictionary array).
    """

    def __init__(self, null_encoding="mask"):
        self._set_options(null_encoding)


cdef class _TakeOptions(FunctionOptions):
    def _set_options(self, boundscheck):
        self.wrapped.reset(new CTakeOptions(boundscheck))


class TakeOptions(_TakeOptions):
    """
    Options for the `take` and `array_take` functions.

    Parameters
    ----------
    boundscheck : boolean, default True
        Whether to check indices are within bounds. If False and an
        index is out of boundes, behavior is undefined (the process
        may crash).
    """

    def __init__(self, *, boundscheck=True):
        self._set_options(boundscheck)


cdef class _MakeStructOptions(FunctionOptions):
    def _set_options(self, field_names, field_nullability, field_metadata):
        cdef:
            vector[c_string] c_field_names
            vector[shared_ptr[const CKeyValueMetadata]] c_field_metadata
        for name in field_names:
            c_field_names.push_back(tobytes(name))
        for metadata in field_metadata:
            c_field_metadata.push_back(pyarrow_unwrap_metadata(metadata))
        self.wrapped.reset(
            new CMakeStructOptions(c_field_names, field_nullability,
                                   c_field_metadata)
        )


class MakeStructOptions(_MakeStructOptions):
    """
    Options for the `make_struct` function.

    Parameters
    ----------
    field_names : sequence of str
        Names of the struct fields to create.
    field_nullability : sequence of bool, optional
        Nullability information for each struct field.
        If omitted, all fields are nullable.
    field_metadata : sequence of KeyValueMetadata, optional
        Metadata for each struct field.
    """

    def __init__(self, field_names=(), *, field_nullability=None,
                 field_metadata=None):
        if field_nullability is None:
            field_nullability = [True] * len(field_names)
        if field_metadata is None:
            field_metadata = [None] * len(field_names)
        self._set_options(field_names, field_nullability, field_metadata)


cdef class _StructFieldOptions(FunctionOptions):
    def _set_options(self, indices):
        cdef:
            CFieldRef field_ref
            const CFieldRef* field_ref_ptr

        if isinstance(indices, (list, tuple)):
            if len(indices):
                indices = Expression._nested_field(tuple(indices))
            else:
                # Allow empty indices; effecitively return same array
                self.wrapped.reset(
                    new CStructFieldOptions(<vector[int]>indices))
                return

        if isinstance(indices, Expression):
            field_ref_ptr = (<Expression>indices).unwrap().field_ref()
            if field_ref_ptr is NULL:
                raise ValueError("Unable to get CFieldRef from Expression")
            field_ref = <CFieldRef>deref(field_ref_ptr)
        elif isinstance(indices, (bytes, str)):
            if indices.startswith(b'.' if isinstance(indices, bytes) else '.'):
                field_ref = GetResultValue(
                    CFieldRef.FromDotPath(<c_string>tobytes(indices)))
            else:
                field_ref = CFieldRef(<c_string>tobytes(indices))
        elif isinstance(indices, int):
            field_ref = CFieldRef(<int> indices)
        else:
            raise TypeError("Expected List[str], List[int], List[bytes], "
                            "Expression, bytes, str, or int. "
                            f"Got: {type(indices)}")
        self.wrapped.reset(new CStructFieldOptions(field_ref))


class StructFieldOptions(_StructFieldOptions):
    """
    Options for the `struct_field` function.

    Parameters
    ----------
    indices : List[str], List[bytes], List[int], Expression, bytes, str, or int
        List of indices for chained field lookup, for example `[4, 1]`
        will look up the second nested field in the fifth outer field.
    """

    def __init__(self, indices):
        self._set_options(indices)


cdef class _ScalarAggregateOptions(FunctionOptions):
    def _set_options(self, skip_nulls, min_count):
        self.wrapped.reset(new CScalarAggregateOptions(skip_nulls, min_count))


class ScalarAggregateOptions(_ScalarAggregateOptions):
    __doc__ = f"""
    Options for scalar aggregations.

    Parameters
    ----------
    {_skip_nulls_doc()}
    {_min_count_doc(default=1)}
    """

    def __init__(self, *, skip_nulls=True, min_count=1):
        self._set_options(skip_nulls, min_count)


cdef class _CountOptions(FunctionOptions):
    _mode_map = {
        "only_valid": CCountMode_ONLY_VALID,
        "only_null": CCountMode_ONLY_NULL,
        "all": CCountMode_ALL,
    }

    def _set_options(self, mode):
        try:
            self.wrapped.reset(new CCountOptions(self._mode_map[mode]))
        except KeyError:
            _raise_invalid_function_option(mode, "count mode")


class CountOptions(_CountOptions):
    """
    Options for the `count` function.

    Parameters
    ----------
    mode : str, default "only_valid"
        Which values to count in the input.
        Accepted values are "only_valid", "only_null", "all".
    """

    def __init__(self, mode="only_valid"):
        self._set_options(mode)


cdef class _IndexOptions(FunctionOptions):
    def _set_options(self, scalar):
        self.wrapped.reset(new CIndexOptions(pyarrow_unwrap_scalar(scalar)))


class IndexOptions(_IndexOptions):
    """
    Options for the `index` function.

    Parameters
    ----------
    value : Scalar
        The value to search for.
    """

    def __init__(self, value):
        self._set_options(value)


cdef class _MapLookupOptions(FunctionOptions):
    _occurrence_map = {
        "all": CMapLookupOccurrence_ALL,
        "first": CMapLookupOccurrence_FIRST,
        "last": CMapLookupOccurrence_LAST,
    }

    def _set_options(self, query_key, occurrence):
        try:
            self.wrapped.reset(
                new CMapLookupOptions(
                    pyarrow_unwrap_scalar(query_key),
                    self._occurrence_map[occurrence]
                )
            )
        except KeyError:
            _raise_invalid_function_option(occurrence,
                                           "Should either be first, last, or all")


class MapLookupOptions(_MapLookupOptions):
    """
    Options for the `map_lookup` function.

    Parameters
    ----------
    query_key : Scalar
        The key to search for.
    occurrence : str
        The occurrence(s) to return from the Map
        Accepted values are "first", "last", or "all".
    """

    def __init__(self, query_key, occurrence):
        self._set_options(query_key, occurrence)


cdef class _ModeOptions(FunctionOptions):
    def _set_options(self, n, skip_nulls, min_count):
        self.wrapped.reset(new CModeOptions(n, skip_nulls, min_count))


class ModeOptions(_ModeOptions):
    __doc__ = f"""
    Options for the `mode` function.

    Parameters
    ----------
    n : int, default 1
        Number of distinct most-common values to return.
    {_skip_nulls_doc()}
    {_min_count_doc(default=0)}
    """

    def __init__(self, n=1, *, skip_nulls=True, min_count=0):
        self._set_options(n, skip_nulls, min_count)


cdef class _SetLookupOptions(FunctionOptions):
    def _set_options(self, value_set, c_bool skip_nulls):
        cdef unique_ptr[CDatum] valset
        if isinstance(value_set, Array):
            valset.reset(new CDatum((<Array> value_set).sp_array))
        elif isinstance(value_set, ChunkedArray):
            valset.reset(
                new CDatum((<ChunkedArray> value_set).sp_chunked_array)
            )
        elif isinstance(value_set, Scalar):
            valset.reset(new CDatum((<Scalar> value_set).unwrap()))
        else:
            _raise_invalid_function_option(value_set, "value set",
                                           exception_class=TypeError)

        self.wrapped.reset(new CSetLookupOptions(deref(valset), skip_nulls))


class SetLookupOptions(_SetLookupOptions):
    """
    Options for the `is_in` and `index_in` functions.

    Parameters
    ----------
    value_set : Array
        Set of values to look for in the input.
    skip_nulls : bool, default False
        If False, nulls in the input are matched in the value_set just
        like regular values.
        If True, nulls in the input always fail matching.
    """

    def __init__(self, value_set, *, skip_nulls=False):
        self._set_options(value_set, skip_nulls)


cdef class _StrptimeOptions(FunctionOptions):
    _unit_map = {
        "s": TimeUnit_SECOND,
        "ms": TimeUnit_MILLI,
        "us": TimeUnit_MICRO,
        "ns": TimeUnit_NANO,
    }

    def _set_options(self, format, unit, error_is_null):
        try:
            self.wrapped.reset(
                new CStrptimeOptions(tobytes(format), self._unit_map[unit],
                                     error_is_null)
            )
        except KeyError:
            _raise_invalid_function_option(unit, "time unit")


class StrptimeOptions(_StrptimeOptions):
    """
    Options for the `strptime` function.

    Parameters
    ----------
    format : str
        Pattern for parsing input strings as timestamps, such as "%Y/%m/%d".
    unit : str
        Timestamp unit of the output.
        Accepted values are "s", "ms", "us", "ns".
    error_is_null : boolean, default False
        Return null on parsing errors if true or raise if false.
    """

    def __init__(self, format, unit, error_is_null=False):
        self._set_options(format, unit, error_is_null)


cdef class _StrftimeOptions(FunctionOptions):
    def _set_options(self, format, locale):
        self.wrapped.reset(
            new CStrftimeOptions(tobytes(format), tobytes(locale))
        )


class StrftimeOptions(_StrftimeOptions):
    """
    Options for the `strftime` function.

    Parameters
    ----------
    format : str, default "%Y-%m-%dT%H:%M:%S"
        Pattern for formatting input values.
    locale : str, default "C"
        Locale to use for locale-specific format specifiers.
    """

    def __init__(self, format="%Y-%m-%dT%H:%M:%S", locale="C"):
        self._set_options(format, locale)


cdef class _DayOfWeekOptions(FunctionOptions):
    def _set_options(self, count_from_zero, week_start):
        self.wrapped.reset(
            new CDayOfWeekOptions(count_from_zero, week_start)
        )


class DayOfWeekOptions(_DayOfWeekOptions):
    """
    Options for the `day_of_week` function.

    Parameters
    ----------
    count_from_zero : bool, default True
        If True, number days from 0, otherwise from 1.
    week_start : int, default 1
        Which day does the week start with (Monday=1, Sunday=7).
        How this value is numbered is unaffected by `count_from_zero`.
    """

    def __init__(self, *, count_from_zero=True, week_start=1):
        self._set_options(count_from_zero, week_start)


cdef class _WeekOptions(FunctionOptions):
    def _set_options(self, week_starts_monday, count_from_zero,
                     first_week_is_fully_in_year):
        self.wrapped.reset(
            new CWeekOptions(week_starts_monday, count_from_zero,
                             first_week_is_fully_in_year)
        )


class WeekOptions(_WeekOptions):
    """
    Options for the `week` function.

    Parameters
    ----------
    week_starts_monday : bool, default True
        If True, weeks start on Monday; if False, on Sunday.
    count_from_zero : bool, default False
        If True, dates at the start of a year that fall into the last week
        of the previous year emit 0.
        If False, they emit 52 or 53 (the week number of the last week
        of the previous year).
    first_week_is_fully_in_year : bool, default False
        If True, week number 0 is fully in January.
        If False, a week that begins on December 29, 30 or 31 is considered
        to be week number 0 of the following year.
    """

    def __init__(self, *, week_starts_monday=True, count_from_zero=False,
                 first_week_is_fully_in_year=False):
        self._set_options(week_starts_monday,
                          count_from_zero, first_week_is_fully_in_year)


cdef class _AssumeTimezoneOptions(FunctionOptions):
    _ambiguous_map = {
        "raise": CAssumeTimezoneAmbiguous_AMBIGUOUS_RAISE,
        "earliest": CAssumeTimezoneAmbiguous_AMBIGUOUS_EARLIEST,
        "latest": CAssumeTimezoneAmbiguous_AMBIGUOUS_LATEST,
    }
    _nonexistent_map = {
        "raise": CAssumeTimezoneNonexistent_NONEXISTENT_RAISE,
        "earliest": CAssumeTimezoneNonexistent_NONEXISTENT_EARLIEST,
        "latest": CAssumeTimezoneNonexistent_NONEXISTENT_LATEST,
    }

    def _set_options(self, timezone, ambiguous, nonexistent):
        if ambiguous not in self._ambiguous_map:
            _raise_invalid_function_option(ambiguous,
                                           "'ambiguous' timestamp handling")
        if nonexistent not in self._nonexistent_map:
            _raise_invalid_function_option(nonexistent,
                                           "'nonexistent' timestamp handling")
        self.wrapped.reset(
            new CAssumeTimezoneOptions(tobytes(timezone),
                                       self._ambiguous_map[ambiguous],
                                       self._nonexistent_map[nonexistent])
        )


class AssumeTimezoneOptions(_AssumeTimezoneOptions):
    """
    Options for the `assume_timezone` function.

    Parameters
    ----------
    timezone : str
        Timezone to assume for the input.
    ambiguous : str, default "raise"
        How to handle timestamps that are ambiguous in the assumed timezone.
        Accepted values are "raise", "earliest", "latest".
    nonexistent : str, default "raise"
        How to handle timestamps that don't exist in the assumed timezone.
        Accepted values are "raise", "earliest", "latest".
    """

    def __init__(self, timezone, *, ambiguous="raise", nonexistent="raise"):
        self._set_options(timezone, ambiguous, nonexistent)


cdef class _NullOptions(FunctionOptions):
    def _set_options(self, nan_is_null):
        self.wrapped.reset(new CNullOptions(nan_is_null))


class NullOptions(_NullOptions):
    """
    Options for the `is_null` function.

    Parameters
    ----------
    nan_is_null : bool, default False
        Whether floating-point NaN values are considered null.
    """

    def __init__(self, *, nan_is_null=False):
        self._set_options(nan_is_null)


cdef class _VarianceOptions(FunctionOptions):
    def _set_options(self, ddof, skip_nulls, min_count):
        self.wrapped.reset(new CVarianceOptions(ddof, skip_nulls, min_count))


class VarianceOptions(_VarianceOptions):
    __doc__ = f"""
    Options for the `variance` and `stddev` functions.

    Parameters
    ----------
    ddof : int, default 0
        Number of degrees of freedom.
    {_skip_nulls_doc()}
    {_min_count_doc(default=0)}
    """

    def __init__(self, *, ddof=0, skip_nulls=True, min_count=0):
        self._set_options(ddof, skip_nulls, min_count)


cdef class _SplitOptions(FunctionOptions):
    def _set_options(self, max_splits, reverse):
        self.wrapped.reset(new CSplitOptions(max_splits, reverse))


class SplitOptions(_SplitOptions):
    """
    Options for splitting on whitespace.

    Parameters
    ----------
    max_splits : int or None, default None
        Maximum number of splits for each input value (unlimited if None).
    reverse : bool, default False
        Whether to start splitting from the end of each input value.
        This only has an effect if `max_splits` is not None.
    """

    def __init__(self, *, max_splits=None, reverse=False):
        if max_splits is None:
            max_splits = -1
        self._set_options(max_splits, reverse)


cdef class _SplitPatternOptions(FunctionOptions):
    def _set_options(self, pattern, max_splits, reverse):
        self.wrapped.reset(
            new CSplitPatternOptions(tobytes(pattern), max_splits, reverse)
        )


class SplitPatternOptions(_SplitPatternOptions):
    """
    Options for splitting on a string pattern.

    Parameters
    ----------
    pattern : str
        String pattern to split on.
    max_splits : int or None, default None
        Maximum number of splits for each input value (unlimited if None).
    reverse : bool, default False
        Whether to start splitting from the end of each input value.
        This only has an effect if `max_splits` is not None.
    """

    def __init__(self, pattern, *, max_splits=None, reverse=False):
        if max_splits is None:
            max_splits = -1
        self._set_options(pattern, max_splits, reverse)


cdef CSortOrder unwrap_sort_order(order) except *:
    if order == "ascending":
        return CSortOrder_Ascending
    elif order == "descending":
        return CSortOrder_Descending
    _raise_invalid_function_option(order, "sort order")


cdef CNullPlacement unwrap_null_placement(null_placement) except *:
    if null_placement == "at_start":
        return CNullPlacement_AtStart
    elif null_placement == "at_end":
        return CNullPlacement_AtEnd
    _raise_invalid_function_option(null_placement, "null placement")


cdef class _PartitionNthOptions(FunctionOptions):
    def _set_options(self, pivot, null_placement):
        self.wrapped.reset(new CPartitionNthOptions(
            pivot, unwrap_null_placement(null_placement)))


class PartitionNthOptions(_PartitionNthOptions):
    """
    Options for the `partition_nth_indices` function.

    Parameters
    ----------
    pivot : int
        Index into the equivalent sorted array of the pivot element.
    null_placement : str, default "at_end"
        Where nulls in the input should be partitioned.
        Accepted values are "at_start", "at_end".
    """

    def __init__(self, pivot, *, null_placement="at_end"):
        self._set_options(pivot, null_placement)


cdef class _CumulativeSumOptions(FunctionOptions):
    def _set_options(self, start, skip_nulls):
        if not isinstance(start, Scalar):
            try:
                start = lib.scalar(start)
            except Exception:
                _raise_invalid_function_option(
                    start, "`start` type for CumulativeSumOptions", TypeError)

        self.wrapped.reset(new CCumulativeSumOptions((<Scalar> start).unwrap(), skip_nulls))


class CumulativeSumOptions(_CumulativeSumOptions):
    """
    Options for `cumulative_sum` function.

    Parameters
    ----------
    start : Scalar, default 0.0
        Starting value for sum computation
    skip_nulls : bool, default False
        When false, the first encountered null is propagated.
    """

    def __init__(self, start=0.0, *, skip_nulls=False):
        self._set_options(start, skip_nulls)


cdef class _ArraySortOptions(FunctionOptions):
    def _set_options(self, order, null_placement):
        self.wrapped.reset(new CArraySortOptions(
            unwrap_sort_order(order), unwrap_null_placement(null_placement)))


class ArraySortOptions(_ArraySortOptions):
    """
    Options for the `array_sort_indices` function.

    Parameters
    ----------
    order : str, default "ascending"
        Which order to sort values in.
        Accepted values are "ascending", "descending".
    null_placement : str, default "at_end"
        Where nulls in the input should be sorted.
        Accepted values are "at_start", "at_end".
    """

    def __init__(self, order="ascending", *, null_placement="at_end"):
        self._set_options(order, null_placement)


cdef class _SortOptions(FunctionOptions):
    def _set_options(self, sort_keys, null_placement):
        cdef vector[CSortKey] c_sort_keys
        for name, order in sort_keys:
            c_sort_keys.push_back(
                CSortKey(tobytes(name), unwrap_sort_order(order))
            )
        self.wrapped.reset(new CSortOptions(
            c_sort_keys, unwrap_null_placement(null_placement)))


class SortOptions(_SortOptions):
    """
    Options for the `sort_indices` function.

    Parameters
    ----------
    sort_keys : sequence of (name, order) tuples
        Names of field/column keys to sort the input on,
        along with the order each field/column is sorted in.
        Accepted values for `order` are "ascending", "descending".
    null_placement : str, default "at_end"
        Where nulls in input should be sorted, only applying to
        columns/fields mentioned in `sort_keys`.
        Accepted values are "at_start", "at_end".
    """

    def __init__(self, sort_keys=(), *, null_placement="at_end"):
        self._set_options(sort_keys, null_placement)


cdef class _SelectKOptions(FunctionOptions):
    def _set_options(self, k, sort_keys):
        cdef vector[CSortKey] c_sort_keys
        for name, order in sort_keys:
            c_sort_keys.push_back(
                CSortKey(tobytes(name), unwrap_sort_order(order))
            )
        self.wrapped.reset(new CSelectKOptions(k, c_sort_keys))


class SelectKOptions(_SelectKOptions):
    """
    Options for top/bottom k-selection.

    Parameters
    ----------
    k : int
        Number of leading values to select in sorted order
        (i.e. the largest values if sort order is "descending",
        the smallest otherwise).
    sort_keys : sequence of (name, order) tuples
        Names of field/column keys to sort the input on,
        along with the order each field/column is sorted in.
        Accepted values for `order` are "ascending", "descending".
    """

    def __init__(self, k, sort_keys):
        self._set_options(k, sort_keys)


cdef class _QuantileOptions(FunctionOptions):
    _interp_map = {
        "linear": CQuantileInterp_LINEAR,
        "lower": CQuantileInterp_LOWER,
        "higher": CQuantileInterp_HIGHER,
        "nearest": CQuantileInterp_NEAREST,
        "midpoint": CQuantileInterp_MIDPOINT,
    }

    def _set_options(self, quantiles, interp, skip_nulls, min_count):
        try:
            self.wrapped.reset(
                new CQuantileOptions(quantiles, self._interp_map[interp],
                                     skip_nulls, min_count)
            )
        except KeyError:
            _raise_invalid_function_option(interp, "quantile interpolation")


class QuantileOptions(_QuantileOptions):
    __doc__ = f"""
    Options for the `quantile` function.

    Parameters
    ----------
    q : double or sequence of double, default 0.5
        Quantiles to compute. All values must be in [0, 1].
    interpolation : str, default "linear"
        How to break ties between competing data points for a given quantile.
        Accepted values are:

        - "linear": compute an interpolation
        - "lower": always use the smallest of the two data points
        - "higher": always use the largest of the two data points
        - "nearest": select the data point that is closest to the quantile
        - "midpoint": compute the (unweighted) mean of the two data points
    {_skip_nulls_doc()}
    {_min_count_doc(default=0)}
    """

    def __init__(self, q=0.5, *, interpolation="linear", skip_nulls=True,
                 min_count=0):
        if not isinstance(q, (list, tuple, np.ndarray)):
            q = [q]
        self._set_options(q, interpolation, skip_nulls, min_count)


cdef class _TDigestOptions(FunctionOptions):
    def _set_options(self, quantiles, delta, buffer_size, skip_nulls,
                     min_count):
        self.wrapped.reset(
            new CTDigestOptions(quantiles, delta, buffer_size, skip_nulls,
                                min_count)
        )


class TDigestOptions(_TDigestOptions):
    __doc__ = f"""
    Options for the `tdigest` function.

    Parameters
    ----------
    q : double or sequence of double, default 0.5
        Quantiles to approximate. All values must be in [0, 1].
    delta : int, default 100
        Compression parameter for the T-digest algorithm.
    buffer_size : int, default 500
        Buffer size for the T-digest algorithm.
    {_skip_nulls_doc()}
    {_min_count_doc(default=0)}
    """

    def __init__(self, q=0.5, *, delta=100, buffer_size=500, skip_nulls=True,
                 min_count=0):
        if not isinstance(q, (list, tuple, np.ndarray)):
            q = [q]
        self._set_options(q, delta, buffer_size, skip_nulls, min_count)


cdef class _Utf8NormalizeOptions(FunctionOptions):
    _form_map = {
        "NFC": CUtf8NormalizeForm_NFC,
        "NFKC": CUtf8NormalizeForm_NFKC,
        "NFD": CUtf8NormalizeForm_NFD,
        "NFKD": CUtf8NormalizeForm_NFKD,
    }

    def _set_options(self, form):
        try:
            self.wrapped.reset(
                new CUtf8NormalizeOptions(self._form_map[form])
            )
        except KeyError:
            _raise_invalid_function_option(form,
                                           "Unicode normalization form")


class Utf8NormalizeOptions(_Utf8NormalizeOptions):
    """
    Options for the `utf8_normalize` function.

    Parameters
    ----------
    form : str
        Unicode normalization form.
        Accepted values are "NFC", "NFKC", "NFD", NFKD".
    """

    def __init__(self, form):
        self._set_options(form)


cdef class _RandomOptions(FunctionOptions):
    def _set_options(self, initializer):
        if initializer == 'system':
            self.wrapped.reset(new CRandomOptions(
                CRandomOptions.FromSystemRandom()))
            return

        if not isinstance(initializer, int):
            try:
                initializer = hash(initializer)
            except TypeError:
                raise TypeError(
                    f"initializer should be 'system', an integer, "
                    f"or a hashable object; got {initializer!r}")

        if initializer < 0:
            initializer += 2**64
        self.wrapped.reset(new CRandomOptions(
            CRandomOptions.FromSeed(initializer)))


class RandomOptions(_RandomOptions):
    """
    Options for random generation.

    Parameters
    ----------
    initializer : int or str
        How to initialize the underlying random generator.
        If an integer is given, it is used as a seed.
        If "system" is given, the random generator is initialized with
        a system-specific source of (hopefully true) randomness.
        Other values are invalid.
    """

    def __init__(self, *, initializer='system'):
        self._set_options(initializer)


cdef class _RankOptions(FunctionOptions):

    _tiebreaker_map = {
        "min": CRankOptionsTiebreaker_Min,
        "max": CRankOptionsTiebreaker_Max,
        "first": CRankOptionsTiebreaker_First,
        "dense": CRankOptionsTiebreaker_Dense,
    }

    def _set_options(self, sort_keys, null_placement, tiebreaker):
        cdef vector[CSortKey] c_sort_keys
        if isinstance(sort_keys, str):
            c_sort_keys.push_back(
                CSortKey(tobytes(""), unwrap_sort_order(sort_keys))
            )
        else:
            for name, order in sort_keys:
                c_sort_keys.push_back(
                    CSortKey(tobytes(name), unwrap_sort_order(order))
                )
        try:
            self.wrapped.reset(
                new CRankOptions(c_sort_keys,
                                 unwrap_null_placement(null_placement),
                                 self._tiebreaker_map[tiebreaker])
            )
        except KeyError:
            _raise_invalid_function_option(tiebreaker, "tiebreaker")


class RankOptions(_RankOptions):
    """
    Options for the `rank` function.

    Parameters
    ----------
    sort_keys : sequence of (name, order) tuples or str, default "ascending"
        Names of field/column keys to sort the input on,
        along with the order each field/column is sorted in.
        Accepted values for `order` are "ascending", "descending".
        Alternatively, one can simply pass "ascending" or "descending" as a string
        if the input is array-like.
    null_placement : str, default "at_end"
        Where nulls in input should be sorted.
        Accepted values are "at_start", "at_end".
    tiebreaker : str, default "first"
        Configure how ties between equal values are handled.
        Accepted values are:

        - "min": Ties get the smallest possible rank in sorted order.
        - "max": Ties get the largest possible rank in sorted order.
        - "first": Ranks are assigned in order of when ties appear in the
                   input. This ensures the ranks are a stable permutation
                   of the input.
        - "dense": The ranks span a dense [1, M] interval where M is the
                   number of distinct values in the input.
    """

    def __init__(self, sort_keys="ascending", *, null_placement="at_end", tiebreaker="first"):
        self._set_options(sort_keys, null_placement, tiebreaker)


def _group_by(args, keys, aggregations):
    cdef:
        vector[CDatum] c_args
        vector[CDatum] c_keys
        vector[CAggregate] c_aggregations
        CDatum result
        CAggregate c_aggr

    _pack_compute_args(args, &c_args)
    _pack_compute_args(keys, &c_keys)

    for aggr_func_name, aggr_opts in aggregations:
        c_aggr.function = tobytes(aggr_func_name)
        if aggr_opts is not None:
            c_aggr.options = (<FunctionOptions?>aggr_opts).wrapped
        else:
            c_aggr.options = <shared_ptr[CFunctionOptions]>nullptr
        c_aggregations.push_back(c_aggr)

    with nogil:
        result = GetResultValue(
            GroupBy(c_args, c_keys, c_aggregations)
        )

    return wrap_datum(result)


cdef class Expression(_Weakrefable):
    """
    A logical expression to be evaluated against some input.

    To create an expression:

    - Use the factory function ``pyarrow.compute.scalar()`` to create a
      scalar (not necessary when combined, see example below).
    - Use the factory function ``pyarrow.compute.field()`` to reference
      a field (column in table).
    - Compare fields and scalars with ``<``, ``<=``, ``==``, ``>=``, ``>``.
    - Combine expressions using python operators ``&`` (logical and),
      ``|`` (logical or) and ``~`` (logical not).
      Note: python keywords ``and``, ``or`` and ``not`` cannot be used
      to combine expressions.
    - Create expression predicates using Expression methods such as
      ``pyarrow.compute.Expression.isin()``.

    Examples
    --------

    >>> import pyarrow.compute as pc
    >>> (pc.field("a") < pc.scalar(3)) | (pc.field("b") > 7)
    <pyarrow.compute.Expression ((a < 3) or (b > 7))>
    >>> pc.field('a') != 3
    <pyarrow.compute.Expression (a != 3)>
    >>> pc.field('a').isin([1, 2, 3])
    <pyarrow.compute.Expression is_in(a, {value_set=int64:[
      1,
      2,
      3
    ], skip_nulls=false})>
    """

    def __init__(self):
        msg = 'Expression is an abstract class thus cannot be initialized.'
        raise TypeError(msg)

    cdef void init(self, const CExpression& sp):
        self.expr = sp

    @staticmethod
    cdef wrap(const CExpression& sp):
        cdef Expression self = Expression.__new__(Expression)
        self.init(sp)
        return self

    cdef inline CExpression unwrap(self):
        return self.expr

    def equals(self, Expression other):
        return self.expr.Equals(other.unwrap())

    def __str__(self):
        return frombytes(self.expr.ToString())

    def __repr__(self):
        return "<pyarrow.compute.{0} {1}>".format(
            self.__class__.__name__, str(self)
        )

    @staticmethod
    def _deserialize(Buffer buffer not None):
        return Expression.wrap(GetResultValue(CDeserializeExpression(
            pyarrow_unwrap_buffer(buffer))))

    def __reduce__(self):
        buffer = pyarrow_wrap_buffer(GetResultValue(
            CSerializeExpression(self.expr)))
        return Expression._deserialize, (buffer,)

    @staticmethod
    cdef Expression _expr_or_scalar(object expr):
        if isinstance(expr, Expression):
            return (<Expression> expr)
        return (<Expression> Expression._scalar(expr))

    @staticmethod
    def _call(str function_name, list arguments, FunctionOptions options=None):
        cdef:
            vector[CExpression] c_arguments
            shared_ptr[CFunctionOptions] c_options

        for argument in arguments:
            if not isinstance(argument, Expression):
                # Attempt to help convert this to an expression
                try:
                    argument = Expression._scalar(argument)
                except ArrowInvalid:
                    raise TypeError(
                        "only other expressions allowed as arguments")
            c_arguments.push_back((<Expression> argument).expr)

        if options is not None:
            c_options = options.unwrap()

        return Expression.wrap(CMakeCallExpression(
            tobytes(function_name), move(c_arguments), c_options))

    def __richcmp__(self, other, int op):
        other = Expression._expr_or_scalar(other)
        return Expression._call({
            Py_EQ: "equal",
            Py_NE: "not_equal",
            Py_GT: "greater",
            Py_GE: "greater_equal",
            Py_LT: "less",
            Py_LE: "less_equal",
        }[op], [self, other])

    def __bool__(self):
        raise ValueError(
            "An Expression cannot be evaluated to python True or False. "
            "If you are using the 'and', 'or' or 'not' operators, use '&', "
            "'|' or '~' instead."
        )

    def __invert__(self):
        return Expression._call("invert", [self])

    def __and__(Expression self, other):
        other = Expression._expr_or_scalar(other)
        return Expression._call("and_kleene", [self, other])

    def __or__(Expression self, other):
        other = Expression._expr_or_scalar(other)
        return Expression._call("or_kleene", [self, other])

    def __add__(Expression self, other):
        other = Expression._expr_or_scalar(other)
        return Expression._call("add_checked", [self, other])

    def __mul__(Expression self, other):
        other = Expression._expr_or_scalar(other)
        return Expression._call("multiply_checked", [self, other])

    def __sub__(Expression self, other):
        other = Expression._expr_or_scalar(other)
        return Expression._call("subtract_checked", [self, other])

    def __truediv__(Expression self, other):
        other = Expression._expr_or_scalar(other)
        return Expression._call("divide_checked", [self, other])

    def is_valid(self):
        """
        Check whether the expression is not-null (valid).

        This creates a new expression equivalent to calling the
        `is_valid` compute function on this expression.

        Returns
        -------
        is_valid : Expression
        """
        return Expression._call("is_valid", [self])

    def is_null(self, bint nan_is_null=False):
        """
        Check whether the expression is null.

        This creates a new expression equivalent to calling the
        `is_null` compute function on this expression.

        Parameters
        ----------
        nan_is_null : boolean, default False
            Whether floating-point NaNs are considered null.

        Returns
        -------
        is_null : Expression
        """
        options = NullOptions(nan_is_null=nan_is_null)
        return Expression._call("is_null", [self], options)

    def cast(self, type=None, safe=None, options=None):
        """
        Explicitly set or change the expression's data type.

        This creates a new expression equivalent to calling the
        `cast` compute function on this expression.

        Parameters
        ----------
        type : DataType, default None
            Type to cast array to.
        safe : boolean, default True
            Whether to check for conversion errors such as overflow.
        options : CastOptions, default None
            Additional checks pass by CastOptions

        Returns
        -------
        cast : Expression
        """
        safe_vars_passed = (safe is not None) or (type is not None)

        if safe_vars_passed and (options is not None):
            raise ValueError("Must either pass values for 'type' and 'safe' or pass a "
                             "value for 'options'")

        if options is None:
            type = ensure_type(type, allow_none=False)
            if safe is False:
                options = CastOptions.unsafe(type)
            else:
                options = CastOptions.safe(type)
        return Expression._call("cast", [self], options)

    def isin(self, values):
        """
        Check whether the expression is contained in values.

        This creates a new expression equivalent to calling the
        `is_in` compute function on this expression.

        Parameters
        ----------
        values : Array or iterable
            The values to check for.

        Returns
        -------
        isin : Expression
            A new expression that, when evaluated, checks whether
            this expression's value is contained in `values`.
        """
        if not isinstance(values, Array):
            values = lib.array(values)

        options = SetLookupOptions(values)
        return Expression._call("is_in", [self], options)

    @staticmethod
    def _field(name_or_idx not None):
        cdef:
            CFieldRef c_field

        if isinstance(name_or_idx, int):
            return Expression.wrap(CMakeFieldExpressionByIndex(name_or_idx))
        else:
            c_field = CFieldRef(<c_string> tobytes(name_or_idx))
            return Expression.wrap(CMakeFieldExpression(c_field))

    @staticmethod
    def _nested_field(tuple names not None):
        cdef:
            vector[CFieldRef] nested

        if len(names) == 0:
            raise ValueError("nested field reference should be non-empty")
        nested.reserve(len(names))
        for name in names:
            if isinstance(name, int):
                nested.push_back(CFieldRef(<int>name))
            else:
                nested.push_back(CFieldRef(<c_string> tobytes(name)))
        return Expression.wrap(CMakeFieldExpression(CFieldRef(move(nested))))

    @staticmethod
    def _scalar(value):
        cdef:
            Scalar scalar

        if isinstance(value, Scalar):
            scalar = value
        else:
            scalar = lib.scalar(value)

        return Expression.wrap(CMakeScalarExpression(scalar.unwrap()))


_deserialize = Expression._deserialize
cdef CExpression _true = CMakeScalarExpression(
    <shared_ptr[CScalar]> make_shared[CBooleanScalar](True)
)


cdef CExpression _bind(Expression filter, Schema schema) except *:
    assert schema is not None

    if filter is None:
        return _true

    return GetResultValue(filter.unwrap().Bind(
        deref(pyarrow_unwrap_schema(schema).get())))


cdef class ScalarUdfContext:
    """
    Per-invocation function context/state.

    This object will always be the first argument to a user-defined
    function. It should not be used outside of a call to the function.
    """

    def __init__(self):
        raise TypeError("Do not call {}'s constructor directly"
                        .format(self.__class__.__name__))

    cdef void init(self, const CScalarUdfContext &c_context):
        self.c_context = c_context

    @property
    def batch_length(self):
        """
        The common length of all input arguments (int).

        In the case that all arguments are scalars, this value
        is used to pass the "actual length" of the arguments,
        e.g. because the scalar values are encoding a column
        with a constant value.
        """
        return self.c_context.batch_length

    @property
    def memory_pool(self):
        """
        A memory pool for allocations (:class:`MemoryPool`).

        This is the memory pool supplied by the user when they invoked
        the function and it should be used in any calls to arrow that the
        UDF makes if that call accepts a memory_pool.
        """
        return box_memory_pool(self.c_context.pool)


cdef inline CFunctionDoc _make_function_doc(dict func_doc) except *:
    """
    Helper function to generate the FunctionDoc
    This function accepts a dictionary and expects the
    summary(str), description(str) and arg_names(List[str]) keys.
    """
    cdef:
        CFunctionDoc f_doc
        vector[c_string] c_arg_names

    f_doc.summary = tobytes(func_doc["summary"])
    f_doc.description = tobytes(func_doc["description"])
    for arg_name in func_doc["arg_names"]:
        c_arg_names.push_back(tobytes(arg_name))
    f_doc.arg_names = c_arg_names
    # UDFOptions integration:
    # TODO: https://issues.apache.org/jira/browse/ARROW-16041
    f_doc.options_class = b""
    f_doc.options_required = False
    return f_doc


cdef object box_scalar_udf_context(const CScalarUdfContext& c_context):
    cdef ScalarUdfContext context = ScalarUdfContext.__new__(ScalarUdfContext)
    context.init(c_context)
    return context


cdef _scalar_udf_callback(user_function, const CScalarUdfContext& c_context, inputs):
    """
    Helper callback function used to wrap the ScalarUdfContext from Python to C++
    execution.
    """
    context = box_scalar_udf_context(c_context)
    return user_function(context, *inputs)


def _get_scalar_udf_context(memory_pool, batch_length):
    cdef CScalarUdfContext c_context
    c_context.pool = maybe_unbox_memory_pool(memory_pool)
    c_context.batch_length = batch_length
    context = box_scalar_udf_context(c_context)
    return context


def register_scalar_function(func, function_name, function_doc, in_types,
                             out_type):
    """
    Register a user-defined scalar function.

    A scalar function is a function that executes elementwise
    operations on arrays or scalars, i.e. a scalar function must
    be computed row-by-row with no state where each output row
    is computed only from its corresponding input row.
    In other words, all argument arrays have the same length,
    and the output array is of the same length as the arguments.
    Scalar functions are the only functions allowed in query engine
    expressions.

    Parameters
    ----------
    func : callable
        A callable implementing the user-defined function.
        The first argument is the context argument of type
        ScalarUdfContext.
        Then, it must take arguments equal to the number of
        in_types defined. It must return an Array or Scalar
        matching the out_type. It must return a Scalar if
        all arguments are scalar, else it must return an Array.

        To define a varargs function, pass a callable that takes
        varargs. The last in_type will be the type of all varargs
        arguments.
    function_name : str
        Name of the function. This name must be globally unique.
    function_doc : dict
        A dictionary object with keys "summary" (str),
        and "description" (str).
    in_types : Dict[str, DataType]
        A dictionary mapping function argument names to
        their respective DataType.
        The argument names will be used to generate
        documentation for the function. The number of
        arguments specified here determines the function
        arity.
    out_type : DataType
        Output type of the function.

    Examples
    --------
    >>> import pyarrow as pa
    >>> import pyarrow.compute as pc
    >>>
    >>> func_doc = {}
    >>> func_doc["summary"] = "simple udf"
    >>> func_doc["description"] = "add a constant to a scalar"
    >>>
    >>> def add_constant(ctx, array):
    ...     return pc.add(array, 1, memory_pool=ctx.memory_pool)
    >>>
    >>> func_name = "py_add_func"
    >>> in_types = {"array": pa.int64()}
    >>> out_type = pa.int64()
    >>> pc.register_scalar_function(add_constant, func_name, func_doc,
    ...                   in_types, out_type)
    >>>
    >>> func = pc.get_function(func_name)
    >>> func.name
    'py_add_func'
    >>> answer = pc.call_function(func_name, [pa.array([20])])
    >>> answer
    <pyarrow.lib.Int64Array object at ...>
    [
      21
    ]
    """
    cdef:
        c_string c_func_name
        CArity c_arity
        CFunctionDoc c_func_doc
        vector[shared_ptr[CDataType]] c_in_types
        PyObject* c_function
        shared_ptr[CDataType] c_out_type
        CScalarUdfOptions c_options

    if callable(func):
        c_function = <PyObject*>func
    else:
        raise TypeError("func must be a callable")

    c_func_name = tobytes(function_name)

    func_spec = inspect.getfullargspec(func)
    num_args = -1
    if isinstance(in_types, dict):
        for in_type in in_types.values():
            c_in_types.push_back(
                pyarrow_unwrap_data_type(ensure_type(in_type)))
        function_doc["arg_names"] = in_types.keys()
        num_args = len(in_types)
    else:
        raise TypeError(
            "in_types must be a dictionary of DataType")

    c_arity = CArity(<int> num_args, func_spec.varargs)

    if "summary" not in function_doc:
        raise ValueError("Function doc must contain a summary")

    if "description" not in function_doc:
        raise ValueError("Function doc must contain a description")

    if "arg_names" not in function_doc:
        raise ValueError("Function doc must contain arg_names")

    c_func_doc = _make_function_doc(function_doc)

    c_out_type = pyarrow_unwrap_data_type(ensure_type(out_type))

    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

    check_status(RegisterScalarFunction(c_function,
                                        <function[CallbackUdf]> &_scalar_udf_callback, c_options))