python/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.

from uuid import UUID

import pyarrow as pa
import pyarrow.compute as pc
import pytest
from datafusion import SessionContext, column, udf
from datafusion import functions as f


@pytest.fixture
def df(ctx):
    # create a RecordBatch and a new DataFrame from it
    batch = pa.RecordBatch.from_arrays(
        [pa.array([1, 2, 3]), pa.array([4, 4, None])],
        names=["a", "b"],
    )
    return ctx.create_dataframe([[batch]], name="test_table")


def test_udf(df):
    # is_null is a pa function over arrays
    is_null = udf(
        lambda x: x.is_null(),
        [pa.int64()],
        pa.bool_(),
        volatility="immutable",
    )

    df = df.select(is_null(column("b")))
    result = df.collect()[0].column(0)

    assert result == pa.array([False, False, True])


def test_udf_decorator(df):
    @udf([pa.int64()], pa.bool_(), "immutable")
    def is_null(x: pa.Array) -> pa.Array:
        return x.is_null()

    df = df.select(is_null(column("b")))
    result = df.collect()[0].column(0)
    assert result == pa.array([False, False, True])


def test_register_udf(ctx, df) -> None:
    is_null = udf(
        lambda x: x.is_null(),
        [pa.float64()],
        pa.bool_(),
        volatility="immutable",
        name="is_null",
    )

    ctx.register_udf(is_null)

    df_result = ctx.sql("select is_null(b) from test_table")
    result = df_result.collect()[0].column(0)

    assert result == pa.array([False, False, True])


def test_udf_lookup(ctx, df) -> None:
    is_null = udf(
        lambda x: x.is_null(),
        [pa.float64()],
        pa.bool_(),
        volatility="immutable",
        name="lookup_is_null",
    )
    ctx.register_udf(is_null)

    assert "lookup_is_null" in ctx.udfs()

    looked_up = ctx.udf("lookup_is_null")
    df_result = df.select(looked_up(column("b")))
    result = df_result.collect()[0].column(0)
    assert result == pa.array([False, False, True])

    with pytest.raises(KeyError, match="no UDF named"):
        ctx.udf("does_not_exist")


def test_udf_late_binding_dispatch(ctx, df) -> None:
    """Resolve a UDF chosen by configuration string, then invoke it."""
    late_is_null = udf(
        lambda x: x.is_null(),
        [pa.int64()],
        pa.bool_(),
        volatility="immutable",
        name="late_is_null",
    )
    late_is_not_null = udf(
        lambda x: pc.invert(x.is_null()),
        [pa.int64()],
        pa.bool_(),
        volatility="immutable",
        name="late_is_not_null",
    )

    ctx.register_udf(late_is_null)
    ctx.register_udf(late_is_not_null)

    # Pretend this came from a config file / API request — only a string.
    runtime_config = {"check_fn": "late_is_not_null"}

    assert runtime_config["check_fn"] in ctx.udfs()

    fn = ctx.udf(runtime_config["check_fn"])
    result = df.select(fn(column("b")).alias("ok")).collect()[0].column(0)
    assert result == pa.array([True, True, False])


def test_udaf_lookup_builtin(ctx, df) -> None:
    assert "sum" in ctx.udafs()
    sum_fn = ctx.udaf("sum")
    result = df.aggregate([], [sum_fn(column("a")).alias("total")]).collect()
    assert result[0].column(0).to_pylist() == [6]

    with pytest.raises(KeyError, match="no UDAF named"):
        ctx.udaf("does_not_exist")


def test_udwf_lookup_builtin(ctx, df) -> None:
    assert "row_number" in ctx.udwfs()
    rn = ctx.udwf("row_number")
    result = df.select(column("a"), rn().alias("rn")).collect()
    assert result[0].column(1).to_pylist() == [1, 2, 3]

    with pytest.raises(KeyError, match="no UDWF named"):
        ctx.udwf("does_not_exist")


class OverThresholdUDF:
    def __init__(self, threshold: int = 0) -> None:
        self.threshold = threshold

    def __call__(self, values: pa.Array) -> pa.Array:
        return pa.array(v.as_py() >= self.threshold for v in values)


def test_udf_with_parameters_function(df) -> None:
    udf_no_param = udf(
        OverThresholdUDF(),
        pa.int64(),
        pa.bool_(),
        volatility="immutable",
    )

    df1 = df.select(udf_no_param(column("a")))
    result = df1.collect()[0].column(0)

    assert result == pa.array([True, True, True])

    udf_with_param = udf(
        OverThresholdUDF(2),
        pa.int64(),
        pa.bool_(),
        volatility="immutable",
    )

    df2 = df.select(udf_with_param(column("a")))
    result = df2.collect()[0].column(0)

    assert result == pa.array([False, True, True])


def test_udf_with_parameters_decorator(df) -> None:
    @udf([pa.int64()], pa.bool_(), "immutable")
    def udf_no_param(values: pa.Array) -> pa.Array:
        return OverThresholdUDF()(values)

    df1 = df.select(udf_no_param(column("a")))
    result = df1.collect()[0].column(0)

    assert result == pa.array([True, True, True])

    @udf([pa.int64()], pa.bool_(), "immutable")
    def udf_with_param(values: pa.Array) -> pa.Array:
        return OverThresholdUDF(2)(values)

    df2 = df.select(udf_with_param(column("a")))
    result = df2.collect()[0].column(0)

    assert result == pa.array([False, True, True])


def test_udf_with_metadata(ctx) -> None:
    @udf([pa.string()], pa.uuid(), "stable")
    def uuid_from_string(uuid_string):
        return pa.array((UUID(s).bytes for s in uuid_string.to_pylist()), pa.uuid())

    @udf([pa.uuid()], pa.int64(), "stable")
    def uuid_version(uuid):
        return pa.array(s.version for s in uuid.to_pylist())

    batch = pa.record_batch({"idx": pa.array(range(5))})
    results = (
        ctx.create_dataframe([[batch]])
        .with_column("uuid_string", f.uuid())
        .with_column("uuid", uuid_from_string(column("uuid_string")))
        .select(uuid_version(column("uuid").alias("uuid_version")))
        .collect()
    )

    assert results[0][0].to_pylist() == [4, 4, 4, 4, 4]


def test_udf_with_nullability(ctx: SessionContext) -> None:
    field_nullable_i64 = pa.field("with_nulls", type=pa.int64(), nullable=True)
    field_non_nullable_i64 = pa.field("no_nulls", type=pa.int64(), nullable=False)

    @udf([field_nullable_i64], field_nullable_i64, "stable")
    def nullable_abs(input_col):
        return pc.abs(input_col)

    @udf([field_non_nullable_i64], field_non_nullable_i64, "stable")
    def non_nullable_abs(input_col):
        return pc.abs(input_col)

    batch = pa.record_batch(
        {
            "with_nulls": pa.array([-2, None, 0, 1, 2]),
            "no_nulls": pa.array([-2, -1, 0, 1, 2]),
        },
        schema=pa.schema(
            [
                field_nullable_i64,
                field_non_nullable_i64,
            ]
        ),
    )
    ctx.register_record_batches("t", [[batch]])
    df = ctx.table("t")

    # Input matches expected, nullable
    df_result = df.select(nullable_abs(column("with_nulls")))
    returned_field = df_result.schema().field(0)
    assert returned_field.nullable
    results = df_result.collect()
    assert results[0][0].to_pylist() == [2, None, 0, 1, 2]

    # Input coercible to expected, nullable
    df_result = df.select(nullable_abs(column("no_nulls")))
    returned_field = df_result.schema().field(0)
    assert returned_field.nullable
    results = df_result.collect()
    assert results[0][0].to_pylist() == [2, 1, 0, 1, 2]

    # Input matches expected, no nulls
    df_result = df.select(non_nullable_abs(column("no_nulls")))
    returned_field = df_result.schema().field(0)
    assert not returned_field.nullable
    results = df_result.collect()
    assert results[0][0].to_pylist() == [2, 1, 0, 1, 2]

    # Invalid - requires non-nullable input but that is not possible
    df_result = df.select(non_nullable_abs(column("with_nulls")))
    returned_field = df_result.schema().field(0)
    assert not returned_field.nullable

    with pytest.raises(Exception) as e_info:
        _results = df_result.collect()
    assert "Invalid argument error" in str(e_info)