venv/lib/python3.11/site-packages/pyarrow/tests/test_udf.py

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


import pytest

try:
    import numpy as np
except ImportError:
    np = None

import pyarrow as pa
from pyarrow import compute as pc

# UDFs are all tested with a dataset scan
pytestmark = pytest.mark.dataset

# For convenience, most of the test here doesn't care about udf func docs
empty_udf_doc = {"summary": "", "description": ""}

try:
    import pyarrow.dataset as ds
except ImportError:
    ds = None


def mock_udf_context(batch_length=10):
    from pyarrow._compute import _get_udf_context
    return _get_udf_context(pa.default_memory_pool(), batch_length)


class MyError(RuntimeError):
    pass


@pytest.fixture(scope="session")
def sum_agg_func_fixture():
    """
    Register a unary aggregate function (mean)
    """
    def func(ctx, x, *args):
        return pa.scalar(np.nansum(x))

    func_name = "sum_udf"
    func_doc = empty_udf_doc

    pc.register_aggregate_function(func,
                                   func_name,
                                   func_doc,
                                   {
                                       "x": pa.float64(),
                                   },
                                   pa.float64()
                                   )
    return func, func_name


@pytest.fixture(scope="session")
def exception_agg_func_fixture():
    def func(ctx, x):
        raise RuntimeError("Oops")
        return pa.scalar(len(x))

    func_name = "y=exception_len(x)"
    func_doc = empty_udf_doc

    pc.register_aggregate_function(func,
                                   func_name,
                                   func_doc,
                                   {
                                       "x": pa.int64(),
                                   },
                                   pa.int64()
                                   )
    return func, func_name


@pytest.fixture(scope="session")
def wrong_output_dtype_agg_func_fixture(scope="session"):
    def func(ctx, x):
        return pa.scalar(len(x), pa.int32())

    func_name = "y=wrong_output_dtype(x)"
    func_doc = empty_udf_doc

    pc.register_aggregate_function(func,
                                   func_name,
                                   func_doc,
                                   {
                                       "x": pa.int64(),
                                   },
                                   pa.int64()
                                   )
    return func, func_name


@pytest.fixture(scope="session")
def wrong_output_type_agg_func_fixture(scope="session"):
    def func(ctx, x):
        return len(x)

    func_name = "y=wrong_output_type(x)"
    func_doc = empty_udf_doc

    pc.register_aggregate_function(func,
                                   func_name,
                                   func_doc,
                                   {
                                       "x": pa.int64(),
                                   },
                                   pa.int64()
                                   )
    return func, func_name


@pytest.fixture(scope="session")
def binary_func_fixture():
    """
    Register a binary scalar function.
    """
    def binary_function(ctx, m, x):
        return pc.call_function("multiply", [m, x],
                                memory_pool=ctx.memory_pool)
    func_name = "y=mx"
    binary_doc = {"summary": "y=mx",
                  "description": "find y from y = mx"}
    pc.register_scalar_function(binary_function,
                                func_name,
                                binary_doc,
                                {"m": pa.int64(),
                                 "x": pa.int64(),
                                 },
                                pa.int64())
    return binary_function, func_name


@pytest.fixture(scope="session")
def ternary_func_fixture():
    """
    Register a ternary scalar function.
    """
    def ternary_function(ctx, m, x, c):
        mx = pc.call_function("multiply", [m, x],
                              memory_pool=ctx.memory_pool)
        return pc.call_function("add", [mx, c],
                                memory_pool=ctx.memory_pool)
    ternary_doc = {"summary": "y=mx+c",
                   "description": "find y from y = mx + c"}
    func_name = "y=mx+c"
    pc.register_scalar_function(ternary_function,
                                func_name,
                                ternary_doc,
                                {
                                    "array1": pa.int64(),
                                    "array2": pa.int64(),
                                    "array3": pa.int64(),
                                },
                                pa.int64())
    return ternary_function, func_name


@pytest.fixture(scope="session")
def varargs_func_fixture():
    """
    Register a varargs scalar function with at least two arguments.
    """
    def varargs_function(ctx, first, *values):
        acc = first
        for val in values:
            acc = pc.call_function("add", [acc, val],
                                   memory_pool=ctx.memory_pool)
        return acc
    func_name = "z=ax+by+c"
    varargs_doc = {"summary": "z=ax+by+c",
                   "description": "find z from z = ax + by + c"
                   }
    pc.register_scalar_function(varargs_function,
                                func_name,
                                varargs_doc,
                                {
                                    "array1": pa.int64(),
                                    "array2": pa.int64(),
                                },
                                pa.int64())
    return varargs_function, func_name


@pytest.fixture(scope="session")
def nullary_func_fixture():
    """
    Register a nullary scalar function.
    """
    def nullary_func(context):
        return pa.array([42] * context.batch_length, type=pa.int64(),
                        memory_pool=context.memory_pool)

    func_doc = {
        "summary": "random function",
        "description": "generates a random value"
    }
    func_name = "test_nullary_func"
    pc.register_scalar_function(nullary_func,
                                func_name,
                                func_doc,
                                {},
                                pa.int64())

    return nullary_func, func_name


@pytest.fixture(scope="session")
def ephemeral_nullary_func_fixture():
    """
    Register a nullary scalar function with an ephemeral Python function.
    This stresses that the Python function object is properly kept alive by the
    registered function.
    """
    def nullary_func(context):
        return pa.array([42] * context.batch_length, type=pa.int64(),
                        memory_pool=context.memory_pool)

    func_doc = {
        "summary": "random function",
        "description": "generates a random value"
    }
    func_name = "test_ephemeral_nullary_func"
    pc.register_scalar_function(nullary_func,
                                func_name,
                                func_doc,
                                {},
                                pa.int64())

    return func_name


@pytest.fixture(scope="session")
def wrong_output_type_func_fixture():
    """
    Register a scalar function which returns something that is neither
    a Arrow scalar or array.
    """
    def wrong_output_type(ctx):
        return 42

    func_name = "test_wrong_output_type"
    in_types = {}
    out_type = pa.int64()
    doc = {
        "summary": "return wrong output type",
        "description": ""
    }
    pc.register_scalar_function(wrong_output_type, func_name, doc,
                                in_types, out_type)
    return wrong_output_type, func_name


@pytest.fixture(scope="session")
def wrong_output_datatype_func_fixture():
    """
    Register a scalar function whose actual output DataType doesn't
    match the declared output DataType.
    """
    def wrong_output_datatype(ctx, array):
        return pc.call_function("add", [array, 1])
    func_name = "test_wrong_output_datatype"
    in_types = {"array": pa.int64()}
    # The actual output DataType will be int64.
    out_type = pa.int16()
    doc = {
        "summary": "return wrong output datatype",
        "description": ""
    }
    pc.register_scalar_function(wrong_output_datatype, func_name, doc,
                                in_types, out_type)
    return wrong_output_datatype, func_name


@pytest.fixture(scope="session")
def wrong_signature_func_fixture():
    """
    Register a scalar function with the wrong signature.
    """
    # Missing the context argument
    def wrong_signature():
        return pa.scalar(1, type=pa.int64())

    func_name = "test_wrong_signature"
    in_types = {}
    out_type = pa.int64()
    doc = {
        "summary": "UDF with wrong signature",
        "description": ""
    }
    pc.register_scalar_function(wrong_signature, func_name, doc,
                                in_types, out_type)
    return wrong_signature, func_name


@pytest.fixture(scope="session")
def raising_func_fixture():
    """
    Register a scalar function which raises a custom exception.
    """
    def raising_func(ctx):
        raise MyError("error raised by scalar UDF")
    func_name = "test_raise"
    doc = {
        "summary": "raising function",
        "description": ""
    }
    pc.register_scalar_function(raising_func, func_name, doc,
                                {}, pa.int64())
    return raising_func, func_name


@pytest.fixture(scope="session")
def unary_vector_func_fixture():
    """
    Register a vector function
    """
    def pct_rank(ctx, x):
        # copy here to get around pandas 1.0 issue
        return pa.array(x.to_pandas().copy().rank(pct=True))

    func_name = "y=pct_rank(x)"
    doc = empty_udf_doc
    pc.register_vector_function(pct_rank, func_name, doc, {
                                'x': pa.float64()}, pa.float64())

    return pct_rank, func_name


@pytest.fixture(scope="session")
def struct_vector_func_fixture():
    """
    Register a vector function that returns a struct array
    """
    def pivot(ctx, k, v, c):
        df = pa.RecordBatch.from_arrays([k, v, c], names=['k', 'v', 'c']).to_pandas()
        df_pivot = df.pivot(columns='c', values='v', index='k').reset_index()
        return pa.RecordBatch.from_pandas(df_pivot).to_struct_array()

    func_name = "y=pivot(x)"
    doc = empty_udf_doc
    pc.register_vector_function(
        pivot, func_name, doc,
        {'k': pa.int64(), 'v': pa.float64(), 'c': pa.utf8()},
        pa.struct([('k', pa.int64()), ('v1', pa.float64()), ('v2', pa.float64())])
    )

    return pivot, func_name


def check_scalar_function(func_fixture,
                          inputs, *,
                          run_in_dataset=True,
                          batch_length=None):
    function, name = func_fixture
    if batch_length is None:
        all_scalar = True
        for arg in inputs:
            if isinstance(arg, pa.Array):
                all_scalar = False
                batch_length = len(arg)
        if all_scalar:
            batch_length = 1

    func = pc.get_function(name)
    assert func.name == name

    result = pc.call_function(name, inputs, length=batch_length)
    expected_output = function(mock_udf_context(batch_length), *inputs)
    assert result == expected_output
    # At the moment there is an issue when handling nullary functions.
    # See: ARROW-15286 and ARROW-16290.
    if run_in_dataset:
        field_names = [f'field{index}' for index, in_arr in inputs]
        table = pa.Table.from_arrays(inputs, field_names)
        dataset = ds.dataset(table)
        func_args = [ds.field(field_name) for field_name in field_names]
        result_table = dataset.to_table(
            columns={'result': ds.field('')._call(name, func_args)})
        assert result_table.column(0).chunks[0] == expected_output


def test_udf_array_unary(unary_func_fixture):
    check_scalar_function(unary_func_fixture,
                          [
                              pa.array([10, 20], pa.int64())
                          ]
                          )


def test_udf_array_binary(binary_func_fixture):
    check_scalar_function(binary_func_fixture,
                          [
                              pa.array([10, 20], pa.int64()),
                              pa.array([2, 4], pa.int64())
                          ]
                          )


def test_udf_array_ternary(ternary_func_fixture):
    check_scalar_function(ternary_func_fixture,
                          [
                              pa.array([10, 20], pa.int64()),
                              pa.array([2, 4], pa.int64()),
                              pa.array([5, 10], pa.int64())
                          ]
                          )


def test_udf_array_varargs(varargs_func_fixture):
    check_scalar_function(varargs_func_fixture,
                          [
                              pa.array([2, 3], pa.int64()),
                              pa.array([10, 20], pa.int64()),
                              pa.array([3, 7], pa.int64()),
                              pa.array([20, 30], pa.int64()),
                              pa.array([5, 10], pa.int64())
                          ]
                          )


def test_registration_errors():
    # validate function name
    doc = {
        "summary": "test udf input",
        "description": "parameters are validated"
    }
    in_types = {"scalar": pa.int64()}
    out_type = pa.int64()

    def test_reg_function(context):
        return pa.array([10])

    with pytest.raises(TypeError):
        pc.register_scalar_function(test_reg_function,
                                    None, doc, in_types,
                                    out_type)

    # validate function
    with pytest.raises(TypeError, match="func must be a callable"):
        pc.register_scalar_function(None, "test_none_function", doc, in_types,
                                    out_type)

    # validate output type
    expected_expr = "DataType expected, got <class 'NoneType'>"
    with pytest.raises(TypeError, match=expected_expr):
        pc.register_scalar_function(test_reg_function,
                                    "test_output_function", doc, in_types,
                                    None)

    # validate input type
    expected_expr = "in_types must be a dictionary of DataType"
    with pytest.raises(TypeError, match=expected_expr):
        pc.register_scalar_function(test_reg_function,
                                    "test_input_function", doc, None,
                                    out_type)

    # register an already registered function
    # first registration
    pc.register_scalar_function(test_reg_function,
                                "test_reg_function", doc, {},
                                out_type)
    # second registration
    expected_expr = "Already have a function registered with name:" \
        + " test_reg_function"
    with pytest.raises(KeyError, match=expected_expr):
        pc.register_scalar_function(test_reg_function,
                                    "test_reg_function", doc, {},
                                    out_type)


def test_varargs_function_validation(varargs_func_fixture):
    _, func_name = varargs_func_fixture

    error_msg = r"VarArgs function 'z=ax\+by\+c' needs at least 2 arguments"

    with pytest.raises(ValueError, match=error_msg):
        pc.call_function(func_name, [42])


def test_function_doc_validation():
    # validate arity
    in_types = {"scalar": pa.int64()}
    out_type = pa.int64()

    # doc with no summary
    func_doc = {
        "description": "desc"
    }

    def add_const(ctx, scalar):
        return pc.call_function("add", [scalar, 1])

    with pytest.raises(ValueError,
                       match="Function doc must contain a summary"):
        pc.register_scalar_function(add_const, "test_no_summary",
                                    func_doc, in_types,
                                    out_type)

    # doc with no description
    func_doc = {
        "summary": "test summary"
    }

    with pytest.raises(ValueError,
                       match="Function doc must contain a description"):
        pc.register_scalar_function(add_const, "test_no_desc",
                                    func_doc, in_types,
                                    out_type)


def test_nullary_function(nullary_func_fixture):
    # XXX the Python compute layer API doesn't let us override batch_length,
    # so only test with the default value of 1.
    check_scalar_function(nullary_func_fixture, [], run_in_dataset=False,
                          batch_length=1)


def test_ephemeral_function(ephemeral_nullary_func_fixture):
    name = ephemeral_nullary_func_fixture
    result = pc.call_function(name, [], length=1)
    assert result.to_pylist() == [42]


def test_wrong_output_type(wrong_output_type_func_fixture):
    _, func_name = wrong_output_type_func_fixture

    with pytest.raises(TypeError,
                       match="Unexpected output type: int"):
        pc.call_function(func_name, [], length=1)


def test_wrong_output_datatype(wrong_output_datatype_func_fixture):
    _, func_name = wrong_output_datatype_func_fixture

    expected_expr = ("Expected output datatype int16, "
                     "but function returned datatype int64")

    with pytest.raises(TypeError, match=expected_expr):
        pc.call_function(func_name, [pa.array([20, 30])])


def test_wrong_signature(wrong_signature_func_fixture):
    _, func_name = wrong_signature_func_fixture

    expected_expr = (r"wrong_signature\(\) takes 0 positional arguments "
                     "but 1 was given")

    with pytest.raises(TypeError, match=expected_expr):
        pc.call_function(func_name, [], length=1)


def test_wrong_datatype_declaration():
    def identity(ctx, val):
        return val

    func_name = "test_wrong_datatype_declaration"
    in_types = {"array": pa.int64()}
    out_type = {}
    doc = {
        "summary": "test output value",
        "description": "test output"
    }
    with pytest.raises(TypeError,
                       match="DataType expected, got <class 'dict'>"):
        pc.register_scalar_function(identity, func_name,
                                    doc, in_types, out_type)


def test_wrong_input_type_declaration():
    def identity(ctx, val):
        return val

    func_name = "test_wrong_input_type_declaration"
    in_types = {"array": None}
    out_type = pa.int64()
    doc = {
        "summary": "test invalid input type",
        "description": "invalid input function"
    }
    with pytest.raises(TypeError,
                       match="DataType expected, got <class 'NoneType'>"):
        pc.register_scalar_function(identity, func_name, doc,
                                    in_types, out_type)


def test_scalar_udf_context(unary_func_fixture):
    # Check the memory_pool argument is properly propagated
    proxy_pool = pa.proxy_memory_pool(pa.default_memory_pool())
    _, func_name = unary_func_fixture

    res = pc.call_function(func_name,
                           [pa.array([1] * 1000, type=pa.int64())],
                           memory_pool=proxy_pool)
    assert res == pa.array([2] * 1000, type=pa.int64())
    assert proxy_pool.bytes_allocated() == 1000 * 8
    # Destroying Python array should destroy underlying C++ memory
    res = None
    assert proxy_pool.bytes_allocated() == 0


def test_raising_func(raising_func_fixture):
    _, func_name = raising_func_fixture
    with pytest.raises(MyError, match="error raised by scalar UDF"):
        pc.call_function(func_name, [], length=1)


def test_scalar_input(unary_func_fixture):
    function, func_name = unary_func_fixture
    res = pc.call_function(func_name, [pa.scalar(10)])
    assert res == pa.scalar(11)


def test_input_lifetime(unary_func_fixture):
    function, func_name = unary_func_fixture

    proxy_pool = pa.proxy_memory_pool(pa.default_memory_pool())
    assert proxy_pool.bytes_allocated() == 0

    v = pa.array([1] * 1000, type=pa.int64(), memory_pool=proxy_pool)
    assert proxy_pool.bytes_allocated() == 1000 * 8
    pc.call_function(func_name, [v])
    assert proxy_pool.bytes_allocated() == 1000 * 8
    # Calling a UDF should not have kept `v` alive longer than required
    v = None
    assert proxy_pool.bytes_allocated() == 0


def _record_batch_from_iters(schema, *iters):
    arrays = [pa.array(list(v), type=schema[i].type)
              for i, v in enumerate(iters)]
    return pa.RecordBatch.from_arrays(arrays=arrays, schema=schema)


def _record_batch_for_range(schema, n):
    return _record_batch_from_iters(schema,
                                    range(n, n + 10),
                                    range(n + 1, n + 11))


def make_udt_func(schema, batch_gen):
    def udf_func(ctx):
        class UDT:
            def __init__(self):
                self.caller = None

            def __call__(self, ctx):
                try:
                    if self.caller is None:
                        self.caller, ctx = batch_gen(ctx).send, None
                    batch = self.caller(ctx)
                except StopIteration:
                    arrays = [pa.array([], type=field.type)
                              for field in schema]
                    batch = pa.RecordBatch.from_arrays(
                        arrays=arrays, schema=schema)
                return batch.to_struct_array()
        return UDT()
    return udf_func


def datasource1_direct():
    """A short dataset"""
    schema = datasource1_schema()

    class Generator:
        def __init__(self):
            self.n = 3

        def __call__(self, ctx):
            if self.n == 0:
                batch = _record_batch_from_iters(schema, [], [])
            else:
                self.n -= 1
                batch = _record_batch_for_range(schema, self.n)
            return batch.to_struct_array()
    return lambda ctx: Generator()


def datasource1_generator():
    schema = datasource1_schema()

    def batch_gen(ctx):
        for n in range(3, 0, -1):
            # ctx =
            yield _record_batch_for_range(schema, n - 1)
    return make_udt_func(schema, batch_gen)


def datasource1_exception():
    schema = datasource1_schema()

    def batch_gen(ctx):
        for n in range(3, 0, -1):
            # ctx =
            yield _record_batch_for_range(schema, n - 1)
        raise RuntimeError("datasource1_exception")
    return make_udt_func(schema, batch_gen)


def datasource1_schema():
    return pa.schema([('', pa.int32()), ('', pa.int32())])


def datasource1_args(func, func_name):
    func_doc = {"summary": f"{func_name} UDT",
                "description": "test {func_name} UDT"}
    in_types = {}
    out_type = pa.struct([("", pa.int32()), ("", pa.int32())])
    return func, func_name, func_doc, in_types, out_type


def _test_datasource1_udt(func_maker):
    schema = datasource1_schema()
    func = func_maker()
    func_name = func_maker.__name__
    func_args = datasource1_args(func, func_name)
    pc.register_tabular_function(*func_args)
    n = 3
    for item in pc.call_tabular_function(func_name):
        n -= 1
        assert item == _record_batch_for_range(schema, n)


def test_udt_datasource1_direct():
    _test_datasource1_udt(datasource1_direct)


def test_udt_datasource1_generator():
    _test_datasource1_udt(datasource1_generator)


def test_udt_datasource1_exception():
    with pytest.raises(RuntimeError, match='datasource1_exception'):
        _test_datasource1_udt(datasource1_exception)


@pytest.mark.numpy
def test_scalar_agg_basic(unary_agg_func_fixture):
    arr = pa.array([10.0, 20.0, 30.0, 40.0, 50.0], pa.float64())
    result = pc.call_function("mean_udf", [arr])
    expected = pa.scalar(30.0)
    assert result == expected


@pytest.mark.numpy
def test_scalar_agg_empty(unary_agg_func_fixture):
    empty = pa.array([], pa.float64())

    with pytest.raises(pa.ArrowInvalid, match='empty inputs'):
        pc.call_function("mean_udf", [empty])


def test_scalar_agg_wrong_output_dtype(wrong_output_dtype_agg_func_fixture):
    arr = pa.array([10, 20, 30, 40, 50], pa.int64())
    with pytest.raises(pa.ArrowTypeError, match="output datatype"):
        pc.call_function("y=wrong_output_dtype(x)", [arr])


def test_scalar_agg_wrong_output_type(wrong_output_type_agg_func_fixture):
    arr = pa.array([10, 20, 30, 40, 50], pa.int64())
    with pytest.raises(pa.ArrowTypeError, match="output type"):
        pc.call_function("y=wrong_output_type(x)", [arr])


@pytest.mark.numpy
def test_scalar_agg_varargs(varargs_agg_func_fixture):
    arr1 = pa.array([10, 20, 30, 40, 50], pa.int64())
    arr2 = pa.array([1.0, 2.0, 3.0, 4.0, 5.0], pa.float64())

    result = pc.call_function(
        "sum_mean", [arr1, arr2]
    )
    expected = pa.scalar(33.0)
    assert result == expected


@pytest.mark.numpy
def test_scalar_agg_exception(exception_agg_func_fixture):
    arr = pa.array([10, 20, 30, 40, 50, 60], pa.int64())

    with pytest.raises(RuntimeError, match='Oops'):
        pc.call_function("y=exception_len(x)", [arr])


@pytest.mark.numpy
def test_hash_agg_basic(unary_agg_func_fixture):
    arr1 = pa.array([10.0, 20.0, 30.0, 40.0, 50.0], pa.float64())
    arr2 = pa.array([4, 2, 1, 2, 1], pa.int32())

    arr3 = pa.array([60.0, 70.0, 80.0, 90.0, 100.0], pa.float64())
    arr4 = pa.array([5, 1, 1, 4, 1], pa.int32())

    table1 = pa.table([arr2, arr1], names=["id", "value"])
    table2 = pa.table([arr4, arr3], names=["id", "value"])
    table = pa.concat_tables([table1, table2])

    result = table.group_by("id").aggregate([("value", "mean_udf")])
    expected = table.group_by("id").aggregate(
        [("value", "mean")]).rename_columns(['id', 'value_mean_udf'])

    assert result.sort_by('id') == expected.sort_by('id')


@pytest.mark.numpy
def test_hash_agg_empty(unary_agg_func_fixture):
    arr1 = pa.array([], pa.float64())
    arr2 = pa.array([], pa.int32())
    table = pa.table([arr2, arr1], names=["id", "value"])

    result = table.group_by("id").aggregate([("value", "mean_udf")])
    expected = pa.table([pa.array([], pa.int32()), pa.array(
        [], pa.float64())], names=['id', 'value_mean_udf'])

    assert result == expected


def test_hash_agg_wrong_output_dtype(wrong_output_dtype_agg_func_fixture):
    arr1 = pa.array([10, 20, 30, 40, 50], pa.int64())
    arr2 = pa.array([4, 2, 1, 2, 1], pa.int32())

    table = pa.table([arr2, arr1], names=["id", "value"])
    with pytest.raises(pa.ArrowTypeError, match="output datatype"):
        table.group_by("id").aggregate([("value", "y=wrong_output_dtype(x)")])


def test_hash_agg_wrong_output_type(wrong_output_type_agg_func_fixture):
    arr1 = pa.array([10, 20, 30, 40, 50], pa.int64())
    arr2 = pa.array([4, 2, 1, 2, 1], pa.int32())
    table = pa.table([arr2, arr1], names=["id", "value"])

    with pytest.raises(pa.ArrowTypeError, match="output type"):
        table.group_by("id").aggregate([("value", "y=wrong_output_type(x)")])


@pytest.mark.numpy
def test_hash_agg_exception(exception_agg_func_fixture):
    arr1 = pa.array([10, 20, 30, 40, 50], pa.int64())
    arr2 = pa.array([4, 2, 1, 2, 1], pa.int32())
    table = pa.table([arr2, arr1], names=["id", "value"])

    with pytest.raises(RuntimeError, match='Oops'):
        table.group_by("id").aggregate([("value", "y=exception_len(x)")])


@pytest.mark.numpy
def test_hash_agg_random(sum_agg_func_fixture):
    """Test hash aggregate udf with randomly sampled data"""

    value_num = 1000000
    group_num = 1000

    arr1 = pa.array(np.repeat(1, value_num), pa.float64())
    arr2 = pa.array(np.random.choice(group_num, value_num), pa.int32())

    table = pa.table([arr2, arr1], names=['id', 'value'])

    result = table.group_by("id").aggregate([("value", "sum_udf")])
    expected = table.group_by("id").aggregate(
        [("value", "sum")]).rename_columns(['id', 'value_sum_udf'])

    assert result.sort_by('id') == expected.sort_by('id')


@pytest.mark.pandas
def test_vector_basic(unary_vector_func_fixture):
    arr = pa.array([10.0, 20.0, 30.0, 40.0, 50.0], pa.float64())
    result = pc.call_function("y=pct_rank(x)", [arr])
    expected = unary_vector_func_fixture[0](None, arr)
    assert result == expected


@pytest.mark.pandas
def test_vector_empty(unary_vector_func_fixture):
    arr = pa.array([1], pa.float64())
    result = pc.call_function("y=pct_rank(x)", [arr])
    expected = unary_vector_func_fixture[0](None, arr)
    assert result == expected


@pytest.mark.pandas
def test_vector_struct(struct_vector_func_fixture):
    k = pa.array(
        [1, 1, 2, 2], pa.int64()
    )
    v = pa.array(
        [1.0, 2.0, 3.0, 4.0], pa.float64()
    )
    c = pa.array(
        ['v1', 'v2', 'v1', 'v2']
    )
    result = pc.call_function("y=pivot(x)", [k, v, c])
    expected = struct_vector_func_fixture[0](None, k, v, c)
    assert result == expected