test_numeric.py

from __future__ import annotations

import itertools
import math
import operator as ops

import numpy as np
import pyarrow as pa
import pytest

from daft import col, lit
from daft.functions import width_bucket
from daft.recordbatch import MicroPartition
from tests.recordbatch import daft_numeric_types

OPS = [
    ops.add,
    ops.sub,
    ops.mul,
    ops.truediv,
    ops.floordiv,
    ops.mod,
    ops.lt,
    ops.le,
    ops.eq,
    ops.ne,
    ops.ge,
    ops.gt,
]


def lists_close_with_nones(a, b):
    if len(a) != len(b):
        return False
    for x, y in zip(a, b):
        if x is None and y is None:
            continue
        if x is not None and y is not None:
            np.testing.assert_allclose([x], [y])
        else:
            return False
    return True


@pytest.mark.parametrize("data_dtype, op", itertools.product(daft_numeric_types, OPS))
def test_table_numeric_expressions(data_dtype, op) -> None:
    a, b = [5, 6, 7, 8], [1, 2, 3, 4]
    pa_table = pa.Table.from_pydict({"a": a, "b": b})

    daft_recordbatch = MicroPartition.from_arrow(pa_table)
    daft_recordbatch = daft_recordbatch.eval_expression_list(
        [op(col("a").cast(data_dtype), col("b").cast(data_dtype)).alias("result")]
    )

    assert len(daft_recordbatch) == 4
    assert daft_recordbatch.column_names() == ["result"]
    pyresult = [op(left, right) for left, right in zip(a, b)]
    assert daft_recordbatch.get_column_by_name("result").to_pylist() == pyresult


@pytest.mark.parametrize("data_dtype, op", itertools.product(daft_numeric_types, OPS))
def test_table_numeric_expressions_with_nulls(data_dtype, op) -> None:
    a, b = [5, 6, None, 8, None], [1, 2, 3, None, None]
    pa_table = pa.Table.from_pydict({"a": a, "b": b})

    daft_recordbatch = MicroPartition.from_arrow(pa_table)
    daft_recordbatch = daft_recordbatch.eval_expression_list(
        [op(col("a").cast(data_dtype), col("b").cast(data_dtype)).alias("result")]
    )

    assert len(daft_recordbatch) == 5
    assert daft_recordbatch.column_names() == ["result"]
    pyresult = [op(left, right) for left, right in zip(a[:2], b[:2])]
    assert daft_recordbatch.get_column_by_name("result").to_pylist()[:2] == pyresult

    assert daft_recordbatch.get_column_by_name("result").to_pylist()[2:] == [None, None, None]


def test_table_numeric_abs() -> None:
    table = MicroPartition.from_pydict({"a": [None, -1.0, 0, 2, 3, None], "b": [-1, -2, 3, 4, None, None]})

    abs_table = table.eval_expression_list([abs(col("a")), col("b").abs()])

    assert [
        abs(v) if v is not None else v for v in table.get_column_by_name("a").to_pylist()
    ] == abs_table.get_column_by_name("a").to_pylist()
    assert [
        abs(v) if v is not None else v for v in table.get_column_by_name("b").to_pylist()
    ] == abs_table.get_column_by_name("b").to_pylist()


def test_table_abs_bad_input() -> None:
    table = MicroPartition.from_pydict({"a": ["a", "b", "c"]})

    with pytest.raises(ValueError, match="Expected input to abs to be numeric"):
        table.eval_expression_list([abs(col("a"))])


def test_table_numeric_ceil() -> None:
    table = MicroPartition.from_pydict(
        {
            "a": [None, -1.0, -0.5, 0, 0.5, 2, None],
            "b": [-1.7, -1.5, -1.3, 0.3, 0.7, None, None],
        }
    )

    ceil_table = table.eval_expression_list([col("a").ceil(), col("b").ceil()])

    assert [
        math.ceil(v) if v is not None else v for v in table.get_column_by_name("a").to_pylist()
    ] == ceil_table.get_column_by_name("a").to_pylist()
    assert [
        math.ceil(v) if v is not None else v for v in table.get_column_by_name("b").to_pylist()
    ] == ceil_table.get_column_by_name("b").to_pylist()


def test_table_ceil_bad_input() -> None:
    table = MicroPartition.from_pydict({"a": ["a", "b", "c"]})

    with pytest.raises(ValueError, match="Expected input to ceil to be numeric"):
        table.eval_expression_list([col("a").ceil()])


def test_table_numeric_floor() -> None:
    table = MicroPartition.from_pydict(
        {
            "a": [None, -1.0, -0.5, 0.0, 0.5, 2, None],
            "b": [-1.7, -1.5, -1.3, 0.3, 0.7, None, None],
        }
    )

    floor_table = table.eval_expression_list([col("a").floor(), col("b").floor()])

    assert [
        math.floor(v) if v is not None else v for v in table.get_column_by_name("a").to_pylist()
    ] == floor_table.get_column_by_name("a").to_pylist()
    assert [
        math.floor(v) if v is not None else v for v in table.get_column_by_name("b").to_pylist()
    ] == floor_table.get_column_by_name("b").to_pylist()


def test_table_floor_bad_input() -> None:
    table = MicroPartition.from_pydict({"a": ["a", "b", "c"]})

    with pytest.raises(ValueError, match="Expected input to floor to be numeric"):
        table.eval_expression_list([col("a").floor()])


def test_table_numeric_sign() -> None:
    table = MicroPartition.from_pydict(
        {
            "a": [None, -1, -5, 0, 5, 2, None],
            "b": [-1.7, -1.5, -1.3, 0.3, 0.7, None, None],
        }
    )
    my_schema = pa.schema([pa.field("uint8", pa.uint8())])
    table_Unsign = MicroPartition.from_arrow(pa.Table.from_arrays([pa.array([None, 0, 1, 2, 3])], schema=my_schema))

    sign_table = table.eval_expression_list([col("a").sign(), col("b").sign()])
    unsign_sign_table = table_Unsign.eval_expression_list([col("uint8").sign()])

    def checkSign(val):
        if val < 0:
            return -1
        if val > 0:
            return 1
        return 0

    assert [
        checkSign(v) if v is not None else v for v in table.get_column_by_name("a").to_pylist()
    ] == sign_table.get_column_by_name("a").to_pylist()
    assert [
        checkSign(v) if v is not None else v for v in table.get_column_by_name("b").to_pylist()
    ] == sign_table.get_column_by_name("b").to_pylist()
    assert [
        checkSign(v) if v is not None else v for v in table_Unsign.get_column_by_name("uint8").to_pylist()
    ] == unsign_sign_table.get_column_by_name("uint8").to_pylist()


def test_table_sign_bad_input() -> None:
    table = MicroPartition.from_pydict({"a": ["a", "b", "c"]})

    with pytest.raises(ValueError, match="Expected input to sign to be numeric"):
        table.eval_expression_list([col("a").sign()])


def test_table_negate_bad_input() -> None:
    table = MicroPartition.from_pydict({"a": ["a", "b", "c"]})
    with pytest.raises(ValueError, match="Expected input to negate to be numeric"):
        table.eval_expression_list([col("a").negate()])


def test_table_numeric_negative() -> None:
    table = MicroPartition.from_pydict(
        {
            "a": [None, -1, -5, 0, 5, 2, None],
            "b": [-1.7, -1.5, -1.3, 0.3, 0.7, None, None],
        }
    )
    sign_table = table.eval_expression_list([col("a").negate(), col("b").negate()])

    # Check signed integers
    a_result = sign_table.to_pydict()["a"]
    expected_a = [None, 1, 5, 0, -5, -2, None]
    assert a_result == expected_a

    # Check floating point
    b_result = sign_table.to_pydict()["b"]
    expected_b = [1.7, 1.5, 1.3, -0.3, -0.7, None, None]
    assert b_result == expected_b


@pytest.mark.parametrize(
    ("fun", "is_arc", "is_co"),
    [
        ("sin", False, False),
        ("cos", False, False),
        ("tan", False, False),
        ("sinh", False, False),
        ("cosh", False, False),
        ("arcsin", True, False),
        ("arccos", True, False),
        ("arctan", True, False),
        ("radians", False, False),
        ("degrees", False, False),
        ("csc", False, True),
        ("sec", False, True),
        ("cot", False, True),
    ],
)
def test_table_numeric_trigonometry(fun: str, is_arc: bool, is_co: bool) -> None:
    np_fun = fun
    if not is_arc:
        table = MicroPartition.from_pydict({"a": [0.0, math.pi, math.pi / 2, math.nan]})
        if is_co:
            reciprocal = {"cot": "tan", "csc": "sin", "sec": "cos"}
            np_fun = reciprocal[fun]
    else:
        table = MicroPartition.from_pydict({"a": [0.0, 1, 0.5, math.nan]})
    s = table.to_pandas()["a"]
    np_result = getattr(np, np_fun)(s)
    if is_co:
        np_result = 1 / np_result

    trigonometry_table = table.eval_expression_list([getattr(col("a"), fun)()])
    assert (
        all(
            x == pytest.approx(y, rel=1e-9) or (math.isnan(x) and math.isnan(y))
            for x, y in zip(trigonometry_table.get_column_by_name("a").to_pylist(), np_result.to_list())
        )
        is True
    )


def test_table_numeric_tanh() -> None:
    table = MicroPartition.from_pydict({"a": [0.0, math.pi, math.pi / 2, math.nan]})
    s = table.to_pandas()["a"]
    np_result = np.tanh(s)
    arct = table.eval_expression_list([col("a").tanh()])
    assert (
        all(
            x - y < 1.0e-10
            or (x is None and y is None)
            or (math.isnan(x) and math.isnan(y) or math.isinf(x) and math.isinf(y))
            for x, y in zip(arct.get_column_by_name("a").to_pylist(), np_result.to_list())
        )
        is True
    )


def test_table_numeric_arc_trigonometry_oor() -> None:
    table = MicroPartition.from_pydict({"a": [math.pi, 2]})
    cot_table = table.eval_expression_list([col("a").arcsin(), col("a").arccos().alias("b")])
    assert all(math.isnan(x) for x in cot_table.get_column_by_name("a").to_pylist())
    assert all(math.isnan(x) for x in cot_table.get_column_by_name("b").to_pylist())


def test_table_numeric_cot() -> None:
    table = MicroPartition.from_pydict({"a": [0, None, math.nan]})
    cot_table = table.eval_expression_list([col("a").cot()])
    expected = [math.inf, None, math.nan]
    assert (
        all(
            x == y or (math.isnan(x) and math.isnan(y)) or (math.isinf(x) and math.isinf(y))
            for x, y in zip(cot_table.get_column_by_name("a").to_pylist(), expected)
        )
        is True
    )


def test_table_numeric_atan2() -> None:
    # cartesian product of y and x tables
    table = MicroPartition.from_pydict(
        {
            "y": [0.0, 1.0, 0.5, -0.5, -0.0, math.nan, 0.0, math.nan],
            "x": [0.0, 0.0, 0.5, 0.5, -10.0, math.nan, math.nan, 1.0],
        }
    )
    pds = table.to_pandas()
    np_result = np.arctan2(pds["y"], pds["x"])

    atan2_table = table.eval_expression_list([col("y").arctan2(col("x"))])
    assert (
        all(
            a == b or (a is None and b is None) or (math.isnan(a) and math.isnan(b))
            for a, b in zip(atan2_table.get_column_by_name("y").to_pylist(), np_result.to_list())
        )
        is True
    )


def test_table_numeric_atan2_literals() -> None:
    table = MicroPartition.from_pydict({"y": [0.0, 1.0, -1.0, math.nan]})
    pds = table.to_pandas()
    literals = [0.0, 1.0, -1.0, math.nan]
    # lhs has value, rhs has literal
    for litv in literals:
        np_result = np.arctan2(pds["y"], np.repeat(litv, len(pds)))
        atan2_table = table.eval_expression_list([col("y").arctan2(lit(litv))])
        assert (
            all(
                a == b or (a is None and b is None) or (math.isnan(a) and math.isnan(b))
                for a, b in zip(atan2_table.get_column_by_name("y").to_pylist(), np_result.to_list())
            )
            is True
        )

    # lhs has literal, rhs has value
    for litv in literals:
        np_result = np.arctan2(np.repeat(litv, len(pds)), pds["y"])
        atan2_table = table.eval_expression_list([lit(litv).arctan2(col("y"))])
        assert (
            all(
                a == b or (a is None and b is None) or (math.isnan(a) and math.isnan(b))
                for a, b in zip(atan2_table.get_column_by_name("literal").to_pylist(), np_result.to_list())
            )
            is True
        )


def test_table_numeric_arctanh() -> None:
    table = MicroPartition.from_pydict({"a": [0.0, 0.5, 0.9, -0.9, -0.5, -0.0, 1, -1.3, math.nan]})
    s = table.to_pandas()["a"]
    np_result = np.arctanh(s)

    arct = table.eval_expression_list([col("a").arctanh()])
    assert (
        all(
            x - y < 1.0e-10
            or (x is None and y is None)
            or (math.isnan(x) and math.isnan(y) or math.isinf(x) and math.isinf(y))
            for x, y in zip(arct.get_column_by_name("a").to_pylist(), np_result.to_list())
        )
        is True
    )


def test_table_numeric_arcsinh() -> None:
    table = MicroPartition.from_pydict({"a": [0.0, 1.0, 0.5, -0.5, -0.0, math.nan]})
    s = table.to_pandas()["a"]
    np_result = np.arcsinh(s)

    arcs = table.eval_expression_list([col("a").arcsinh()])
    assert (
        all(
            x - y < 1.0e-10 or (x is None and y is None) or (math.isnan(x) and math.isnan(y))
            for x, y in zip(arcs.get_column_by_name("a").to_pylist(), np_result.to_list())
        )
        is True
    )


def test_table_numeric_arccosh() -> None:
    table = MicroPartition.from_pydict({"a": [1.0, 2.0, 1.5, 0.5, math.nan]})
    s = table.to_pandas()["a"]
    np_result = np.arccosh(s)

    arcc = table.eval_expression_list([col("a").arccosh()])
    assert (
        all(
            x - y < 1.0e-10 or (x is None and y is None) or (math.isnan(x) and math.isnan(y))
            for x, y in zip(arcc.get_column_by_name("a").to_pylist(), np_result.to_list())
        )
        is True
    )


def test_table_numeric_round() -> None:
    from decimal import ROUND_HALF_UP, Decimal

    table = MicroPartition.from_pydict(
        {
            "a": [None, -1, -5, 0, 5, 2, None],
            "b": [-1.765, -1.565, -1.321, 0.399, 0.781, None, None],
        }
    )
    round_table = table.eval_expression_list([col("a").round(0), col("b").round(2)])
    assert [
        Decimal(v).to_integral_value(rounding=ROUND_HALF_UP) if v is not None else v
        for v in table.get_column_by_name("a").to_pylist()
    ] == round_table.get_column_by_name("a").to_pylist()
    assert [
        (float(Decimal(str(v)).quantize(Decimal("0.00"), rounding=ROUND_HALF_UP)) if v is not None else v)
        for v in table.get_column_by_name("b").to_pylist()
    ] == round_table.get_column_by_name("b").to_pylist()


def test_table_round_bad_input() -> None:
    table = MicroPartition.from_pydict({"a": ["a", "b", "c"]})

    with pytest.raises(ValueError, match="Expected input to be numeric"):
        table.eval_expression_list([col("a").round()])

    table = MicroPartition.from_pydict({"a": [1, 2, 3]})

    with pytest.raises(ValueError, match="failed to cast"):
        table.eval_expression_list([col("a").round(-2)])


def test_clip_one_sided_bounding():
    table = MicroPartition.from_pydict({"a": [1, 2, 3, 4, 5]})
    clip_table = table.eval_expression_list([col("a").clip(None, 3)])
    expected = [1, 2, 3, 3, 3]
    assert clip_table.get_column_by_name("a").to_pylist() == expected

    clip_table = table.eval_expression_list([col("a").clip(3, None)])
    expected = [3, 3, 3, 4, 5]
    assert clip_table.get_column_by_name("a").to_pylist() == expected


def test_clip_integer_float_promotion():
    table = MicroPartition.from_pydict({"a": [1, 2, 3, 4, 5]})
    clip_table = table.eval_expression_list([col("a").clip(2.5, 4.5)])
    expected = [2.5, 2.5, 3, 4, 4.5]
    assert clip_table.get_column_by_name("a").to_pylist() == expected


def test_clip_zero_handling():
    table = MicroPartition.from_pydict({"a": [-0.0, 0.0, 1.0, -1.0]})
    clip_table = table.eval_expression_list([col("a").clip(-0.5, 0.5)])
    expected = [-0.0, 0.0, 0.5, -0.5]
    assert clip_table.get_column_by_name("a").to_pylist() == expected


def test_clip_empty_array():
    table = MicroPartition.from_pydict({"a": []})
    with pytest.raises(ValueError):
        table.eval_expression_list([col("a").clip(0, 1)])


def test_clip_all_within_bounds():
    table = MicroPartition.from_pydict({"a": [2, 3, 4]})
    clip_table = table.eval_expression_list([col("a").clip(1, 5)])
    expected = [2, 3, 4]
    assert clip_table.get_column_by_name("a").to_pylist() == expected


def test_clip_all_out_of_bounds():
    table = MicroPartition.from_pydict({"a": [0, 1, 2, 3, 4, 5]})
    clip_table = table.eval_expression_list([col("a").clip(2, 3)])
    expected = [2, 2, 2, 3, 3, 3]
    assert clip_table.get_column_by_name("a").to_pylist() == expected


def test_clip_nan_handling():
    table = MicroPartition.from_pydict({"a": [1, 2, np.nan, 4, 5]})
    clip_table = table.eval_expression_list([col("a").clip(2, 4)])
    expected = [2, 2, np.nan, 4, 4]
    actual = clip_table.get_column_by_name("a").to_pylist()
    assert all((a == b or (np.isnan(a) and np.isnan(b))) for a, b in zip(actual, expected))


@pytest.mark.parametrize(
    "lower_bound, upper_bound, expected",
    [
        # Test case 1: Column lower bound with scalar upper bound
        ("lower_bound", 5, [1.0, 2.5, None, 4.7, 5.0, float("nan")]),
        # Test case 2: Scalar lower bound with column upper bound
        (2.0, "upper_bound", [2.0, 2.5, None, 4.7, 5.0, float("nan")]),
        # Test case 3: Column lower and upper bounds
        ("lower_bound", "upper_bound", [1.0, 2.5, None, 4.7, 5.0, float("nan")]),
        # Test case 4: Infinite bounds
        (float("-inf"), float("inf"), [1.0, 2.5, None, 4.7, 5.0, float("nan")]),
        # Test case 5: None bounds
        (None, None, [1.0, 2.5, None, 4.7, 5.0, float("nan")]),
        # Test case 6: Scalar bounds
        (2.0, 5.0, [2.0, 2.5, None, 4.7, 5.0, float("nan")]),
    ],
)
def test_clip(lower_bound, upper_bound, expected):
    """Test clipping a column with various combinations of scalar and column bounds."""
    # Initialize test data
    table = MicroPartition.from_pydict(
        {
            "data": [1.0, 2.5, None, 4.7, 5.0, float("nan")],
            "lower_bound": [0.5, 2.0, 1.0, None, 4.0, 0.0],
            "upper_bound": [2.0, 3.0, 5.0, None, None, float("inf")],
        }
    )

    # Prepare the clip expression
    lower = col(lower_bound) if isinstance(lower_bound, str) else lower_bound
    upper = col(upper_bound) if isinstance(upper_bound, str) else upper_bound

    # Perform the clip operation
    clip_table = table.eval_expression_list([col("data").clip(lower, upper)])
    actual = clip_table.get_column_by_name("data").to_pylist()

    # Verify results
    assert all(
        (a == b) or (a is None and b is None) or (math.isnan(a) and math.isnan(b)) for a, b in zip(actual, expected)
    ), f"Expected {expected}, got {actual}"


def test_clip_incompatible_lengths():
    table1 = MicroPartition.from_pydict({"data": [1, 2, 3, 4, 5]})
    table2 = MicroPartition.from_pydict({"data": [1, 2, 3]})  # Shorter array
    table3 = MicroPartition.from_pydict({"data": [1, 2, 3, 4, 5, 6]})  # Longer array

    # Test shorter lower bound
    with pytest.raises(ValueError):
        table1.eval_expression_list([col("data").clip(table2.get_column_by_name("data"), lit(5))])

    # Test longer upper bound
    with pytest.raises(ValueError):
        table1.eval_expression_list([col("data").clip(lit(0), table3.get_column_by_name("data"))])

    # Test both bounds with different lengths
    with pytest.raises(ValueError):
        table1.eval_expression_list(
            [col("data").clip(table2.get_column_by_name("data"), table3.get_column_by_name("data"))]
        )


def test_clip_unsupported_types():
    # Test with string data
    table = MicroPartition.from_pydict({"strings": ["a", "b", "c"], "numbers": [1, 2, 3]})

    # Test string column as data
    with pytest.raises(ValueError):
        table.eval_expression_list([col("strings").clip(0, 5)])

    # Test string column as lower bound
    with pytest.raises(ValueError):
        table.eval_expression_list([col("numbers").clip(col("strings"), 5)])

    # Test string column as upper bound
    with pytest.raises(ValueError):
        table.eval_expression_list([col("numbers").clip(0, col("strings"))])

    # Test string literals as bounds
    with pytest.raises(ValueError):
        table.eval_expression_list([col("numbers").clip("a", "z")])


def test_table_numeric_log2() -> None:
    table = MicroPartition.from_pydict({"a": [0.1, 0.01, 1.5, None], "b": [1, 10, None, None]})
    log2_table = table.eval_expression_list([col("a").log2(), col("b").log2()])
    assert [
        math.log2(v) if v is not None else v for v in table.get_column_by_name("a").to_pylist()
    ] == log2_table.get_column_by_name("a").to_pylist()
    assert [
        math.log2(v) if v is not None else v for v in table.get_column_by_name("b").to_pylist()
    ] == log2_table.get_column_by_name("b").to_pylist()


def test_table_log2_bad_input() -> None:
    table = MicroPartition.from_pydict({"a": ["a", "b", "c"]})

    with pytest.raises(ValueError, match="Expected input to log to be numeric"):
        table.eval_expression_list([col("a").log2()])


def test_table_numeric_log10() -> None:
    table = MicroPartition.from_pydict({"a": [0.1, 0.01, 1.5, None], "b": [1, 10, None, None]})
    log10_table = table.eval_expression_list([col("a").log10(), col("b").log10()])
    assert [
        math.log10(v) if v is not None else v for v in table.get_column_by_name("a").to_pylist()
    ] == log10_table.get_column_by_name("a").to_pylist()
    assert [
        math.log10(v) if v is not None else v for v in table.get_column_by_name("b").to_pylist()
    ] == log10_table.get_column_by_name("b").to_pylist()


def test_table_log10_bad_input() -> None:
    table = MicroPartition.from_pydict({"a": ["a", "b", "c"]})

    with pytest.raises(ValueError, match="Expected input to log to be numeric"):
        table.eval_expression_list([col("a").log10()])


@pytest.mark.parametrize(
    ("base"),
    [2, 10, 100, math.e],
)
def test_table_numeric_log(base: float) -> None:
    table = MicroPartition.from_pydict({"a": [0.1, 0.01, 1.5, None], "b": [1, 10, None, None]})
    log_table = table.eval_expression_list([col("a").log(base), col("b").log(base)])

    assert lists_close_with_nones(
        log_table.get_column_by_name("a").to_pylist(),
        [math.log(v, base) if v is not None else None for v in table.get_column_by_name("a").to_pylist()],
    )

    assert lists_close_with_nones(
        log_table.get_column_by_name("b").to_pylist(),
        [math.log(v, base) if v is not None else None for v in table.get_column_by_name("b").to_pylist()],
    )


def test_table_log_bad_input() -> None:
    table = MicroPartition.from_pydict({"a": ["a", "b", "c"]})

    with pytest.raises(ValueError, match="Expected input to log to be numeric"):
        table.eval_expression_list([col("a").log(base=10)])


def test_table_numeric_ln() -> None:
    table = MicroPartition.from_pydict({"a": [0.1, 0.01, 1.5, None], "b": [1, 10, None, None]})
    ln_table = table.eval_expression_list([col("a").ln(), col("b").ln()])
    assert lists_close_with_nones(
        [math.log(v) if v is not None else v for v in table.get_column_by_name("a").to_pylist()],
        ln_table.get_column_by_name("a").to_pylist(),
    )
    assert lists_close_with_nones(
        [math.log(v) if v is not None else v for v in table.get_column_by_name("b").to_pylist()],
        ln_table.get_column_by_name("b").to_pylist(),
    )


def test_table_ln_bad_input() -> None:
    table = MicroPartition.from_pydict({"a": ["a", "b", "c"]})

    with pytest.raises(ValueError, match="Expected input to log to be numeric"):
        table.eval_expression_list([col("a").ln()])


def test_table_numeric_log1p() -> None:
    table = MicroPartition.from_pydict({"a": [0.1, 0.01, 1.5, None], "b": [1, 10, None, None]})
    log1p_table = table.eval_expression_list([col("a").log1p(), col("b").log1p()])
    assert lists_close_with_nones(
        [math.log1p(v) if v is not None else v for v in table.get_column_by_name("a").to_pylist()],
        log1p_table.get_column_by_name("a").to_pylist(),
    )
    assert lists_close_with_nones(
        [math.log1p(v) if v is not None else v for v in table.get_column_by_name("b").to_pylist()],
        log1p_table.get_column_by_name("b").to_pylist(),
    )


def test_table_log1p_bad_input() -> None:
    table = MicroPartition.from_pydict({"a": ["a", "b", "c"]})

    with pytest.raises(ValueError, match="Expected input to log to be numeric"):
        table.eval_expression_list([col("a").log1p()])


def test_table_numeric_pow() -> None:
    table = MicroPartition.from_pydict({"a": [0.1, 0.01, 1.5, 3, None], "b": [1, 10, 100, None, None]})
    exprs = [0.01, 0.1, 1, 2, 10]
    for expr in exprs:
        pow_table = table.eval_expression_list([col("a").pow(expr), col("b").pow(expr)])
        assert lists_close_with_nones(
            [math.pow(v, expr) if v is not None else v for v in table.get_column_by_name("a").to_pylist()],
            pow_table.get_column_by_name("a").to_pylist(),
        )
        assert lists_close_with_nones(
            [math.pow(v, expr) if v is not None else v for v in table.get_column_by_name("b").to_pylist()],
            pow_table.get_column_by_name("b").to_pylist(),
        )


def test_table_pow_bad_input() -> None:
    table = MicroPartition.from_pydict({"a": ["a", "b", "c"]})
    table2 = MicroPartition.from_pydict({"a": [1, 10, 100]})

    with pytest.raises(ValueError, match="Expected input to compute pow to be numeric, got String"):
        table.eval_expression_list([col("a").pow(1)])

    with pytest.raises(ValueError, match="Expected input to compute pow to be numeric, got String"):
        table.eval_expression_list([col("a").pow(0.1)])

    with pytest.raises(ValueError, match='DaftError::ValueError Expected floating point number, received: "c"'):
        table2.eval_expression_list([col("a").pow("c")])


def test_table_numeric_power() -> None:
    table = MicroPartition.from_pydict({"a": [0.1, 0.01, 1.5, 3, None], "b": [1, 10, 100, None, None]})
    exprs = [0.01, 0.1, 1, 2, 10]
    for expr in exprs:
        power_table = table.eval_expression_list([col("a").power(expr), col("b").power(expr)])
        assert lists_close_with_nones(
            [math.pow(v, expr) if v is not None else v for v in table.get_column_by_name("a").to_pylist()],
            power_table.get_column_by_name("a").to_pylist(),
        )
        assert lists_close_with_nones(
            [math.pow(v, expr) if v is not None else v for v in table.get_column_by_name("b").to_pylist()],
            power_table.get_column_by_name("b").to_pylist(),
        )


def test_table_power_bad_input() -> None:
    table = MicroPartition.from_pydict({"a": ["a", "b", "c"]})
    table2 = MicroPartition.from_pydict({"a": [1, 10, 100]})

    with pytest.raises(ValueError, match="Expected input to compute power to be numeric, got String"):
        table.eval_expression_list([col("a").power(1)])

    with pytest.raises(ValueError, match="Expected input to compute power to be numeric, got String"):
        table.eval_expression_list([col("a").power(0.1)])

    with pytest.raises(ValueError, match='DaftError::ValueError Expected floating point number, received: "c"'):
        table2.eval_expression_list([col("a").power("c")])


def test_table_expm1() -> None:
    table = MicroPartition.from_pydict({"a": [0.1, 0.01, None], "b": [1, 10, None]})
    expm1_table = table.eval_expression_list([col("a").expm1(), col("b").expm1()])
    assert lists_close_with_nones(
        [0.10517091807564763, 0.010050167084168058, None],
        expm1_table.get_column_by_name("a").to_pylist(),
    )
    assert lists_close_with_nones(
        [1.718281828459045, 22025.465794806718, None],
        expm1_table.get_column_by_name("b").to_pylist(),
    )


def test_table_exp() -> None:
    table = MicroPartition.from_pydict({"a": [0.1, 0.01, None], "b": [1, 10, None]})
    exp_table = table.eval_expression_list([col("a").exp(), col("b").exp()])
    assert lists_close_with_nones(
        [1.1051709180756477, 1.010050167084168, None],
        exp_table.get_column_by_name("a").to_pylist(),
    )
    assert lists_close_with_nones(
        [2.718281828459045, 22026.465794806718, None],
        exp_table.get_column_by_name("b").to_pylist(),
    )


def test_table_numeric_sqrt() -> None:
    table = MicroPartition.from_pydict({"a": [4, 9, None, 16, 25, None], "b": [2.25, 0.81, None, 1, 10.24, None]})
    sqrt_table = table.eval_expression_list([col("a").sqrt(), col("b").sqrt()])
    assert [
        math.sqrt(v) if v is not None else v for v in table.get_column_by_name("a").to_pylist()
    ] == sqrt_table.get_column_by_name("a").to_pylist()
    assert [
        math.sqrt(v) if v is not None else v for v in table.get_column_by_name("b").to_pylist()
    ] == sqrt_table.get_column_by_name("b").to_pylist()


@pytest.mark.parametrize(
    "left, right",
    [
        pytest.param([1, 2, 3], [4, 5, 6], id="NoNulls"),
        pytest.param([1, 2, None], [4, None, 6], id="WithNulls"),
    ],
)
def test_table_shift_left(left, right) -> None:
    table = MicroPartition.from_pydict({"a": left, "b": right})
    shift_left_table = table.eval_expression_list([col("a") << (col("b"))])
    assert shift_left_table.get_column_by_name("a").to_pylist() == [
        i << j if i is not None and j is not None else None for i, j in zip(left, right)
    ]


def test_table_shift_left_with_scalar() -> None:
    table = MicroPartition.from_pydict({"a": [1, 2, 4]})
    shift_left_table = table.eval_expression_list([col("a") << (1)])
    assert [1 << 1, 2 << 1, 4 << 1] == shift_left_table.get_column_by_name("a").to_pylist()


def test_table_shift_left_bad_input() -> None:
    table = MicroPartition.from_pydict({"a": ["a", "b", "c"]})

    with pytest.raises(ValueError, match="Cannot operate shift left on types: String, String"):
        table.eval_expression_list([col("a") << (col("a"))])


def test_table_shift_left_bad_shift() -> None:
    table = MicroPartition.from_pydict({"a": [1, 2, 4], "b": [3, 2, 1]})

    with pytest.raises(ValueError, match="Cannot operate shift left on types: Int64, String"):
        table.eval_expression_list([col("a") << (lit("a"))])


def test_table_shift_left_negative_bits() -> None:
    table = MicroPartition.from_pydict({"a": [1, 2, 4], "b": [3, 2, -1]})

    with pytest.raises(ValueError, match="Cannot shift left by a negative number"):
        table.eval_expression_list([col("a") << (col("b"))])


def test_table_shift_left_syntactic_sugar() -> None:
    table = MicroPartition.from_pydict({"a": [1, 2, 4]})
    shift_table = table.eval_expression_list([col("a").shift_left(1)])
    assert [1 << 1, 2 << 1, 4 << 1] == shift_table.get_column_by_name("a").to_pylist()


@pytest.mark.parametrize(
    "left, right",
    [
        pytest.param([1, 2, 3], [4, 5, 6], id="NoNulls"),
        pytest.param([1, 2, None], [4, None, 6], id="WithNulls"),
    ],
)
def test_table_shift_right(left, right) -> None:
    table = MicroPartition.from_pydict({"a": left, "b": right})
    shift_right_table = table.eval_expression_list([col("a") >> (col("b"))])
    assert shift_right_table.get_column_by_name("a").to_pylist() == [
        i >> j if i is not None and j is not None else None for i, j in zip(left, right)
    ]


def test_table_shift_right_with_scalar() -> None:
    table = MicroPartition.from_pydict({"a": [1, 2, 4]})
    shift_right_table = table.eval_expression_list([col("a") >> (1)])
    assert [1 >> 1, 2 >> 1, 4 >> 1] == shift_right_table.get_column_by_name("a").to_pylist()


def test_table_shift_right_bad_input() -> None:
    table = MicroPartition.from_pydict({"a": ["a", "b", "c"]})

    with pytest.raises(ValueError, match="Cannot operate shift right on types: String, String"):
        table.eval_expression_list([col("a") >> (col("a"))])


def test_table_shift_right_bad_shift() -> None:
    table = MicroPartition.from_pydict({"a": [8, 4, 2], "b": [3, 2, 1]})

    with pytest.raises(ValueError, match="Cannot operate shift right on types: Int64, String"):
        table.eval_expression_list([col("a") >> (lit("a"))])


def test_table_shift_right_negative_bits() -> None:
    table = MicroPartition.from_pydict({"a": [8, 4, 2], "b": [3, 2, -1]})

    with pytest.raises(ValueError, match="Cannot shift right by a negative number"):
        table.eval_expression_list([col("a") >> (col("b"))])


def test_table_shift_right_syntactic_sugar() -> None:
    table = MicroPartition.from_pydict({"a": [1, 2, 4]})
    shift_table = table.eval_expression_list([col("a").shift_right(1)])
    assert [1 >> 1, 2 >> 1, 4 >> 1] == shift_table.get_column_by_name("a").to_pylist()


def test_e() -> None:
    from daft.functions import e

    table = MicroPartition.from_pydict({"a": [1, 2, 3]})
    result = table.eval_expression_list([e().alias("e")])
    values = result.get_column_by_name("e").to_pylist()
    assert len(values) == 3
    for v in values:
        assert v == pytest.approx(math.e, rel=1e-10)


def test_pi() -> None:
    from daft.functions import pi

    table = MicroPartition.from_pydict({"a": [1, 2, 3]})
    result = table.eval_expression_list([pi().alias("pi")])
    values = result.get_column_by_name("pi").to_pylist()
    assert len(values) == 3
    for v in values:
        assert v == pytest.approx(math.pi, rel=1e-10)


def test_factorial() -> None:
    from daft.functions import factorial

    table = MicroPartition.from_pydict({"a": [0, 1, 2, 3, 4, 5, 10, None]})
    result = table.eval_expression_list([factorial(col("a")).alias("result")])
    values = result.get_column_by_name("result").to_pylist()
    expected = [1, 1, 2, 6, 24, 120, 3628800, None]
    assert values == expected


def test_factorial_negative() -> None:
    from daft.functions import factorial

    table = MicroPartition.from_pydict({"a": [-1, -5]})
    result = table.eval_expression_list([factorial(col("a")).alias("result")])
    values = result.get_column_by_name("result").to_pylist()
    assert values == [None, None]


def test_factorial_overflow() -> None:
    from daft.functions import factorial

    table = MicroPartition.from_pydict({"a": [20, 21, 100]})
    result = table.eval_expression_list([factorial(col("a")).alias("result")])
    values = result.get_column_by_name("result").to_pylist()
    assert values == [2432902008176640000, None, None]


def test_factorial_float_input() -> None:
    from daft.functions import factorial

    table = MicroPartition.from_pydict({"a": [3.5, 2.1]})
    with pytest.raises(ValueError, match="Expected input to factorial to be integer"):
        table.eval_expression_list([factorial(col("a"))])


def test_factorial_bad_input() -> None:
    from daft.functions import factorial

    table = MicroPartition.from_pydict({"a": ["a", "b", "c"]})
    with pytest.raises(ValueError, match="Expected input to factorial to be integer"):
        table.eval_expression_list([factorial(col("a"))])


def test_bin() -> None:
    from daft.functions import bin

    table = MicroPartition.from_pydict({"a": [0, 1, 2, 7, 255, 1024, None]})
    result = table.eval_expression_list([bin(col("a")).alias("result")])
    values = result.get_column_by_name("result").to_pylist()
    expected = ["0", "1", "10", "111", "11111111", "10000000000", None]
    assert values == expected


def test_bin_negative() -> None:
    from daft.functions import bin

    table = MicroPartition.from_pydict({"a": pa.array([-1, -2, -(2**63)], type=pa.int64())})
    result = table.eval_expression_list([bin(col("a")).alias("result")])
    values = result.get_column_by_name("result").to_pylist()
    assert values == ["1" * 64, "1" * 63 + "0", "1" + "0" * 63]


def test_bin_unsigned_input() -> None:
    from daft.functions import bin

    table = MicroPartition.from_pydict({"a": pa.array([0, 7, 255], type=pa.uint16())})
    result = table.eval_expression_list([bin(col("a")).alias("result")])
    values = result.get_column_by_name("result").to_pylist()
    assert values == ["0", "111", "11111111"]


def test_bin_signed_small_width() -> None:
    from daft.functions import bin

    table = MicroPartition.from_pydict({"a": pa.array([-1, -128], type=pa.int8())})
    result = table.eval_expression_list([bin(col("a")).alias("result")])
    values = result.get_column_by_name("result").to_pylist()
    assert values == ["1" * 64, "1" * 56 + "10000000"]


def test_bin_uint64_max() -> None:
    from daft.functions import bin

    table = MicroPartition.from_pydict({"a": pa.array([0, 2**63, 2**64 - 1], type=pa.uint64())})
    result = table.eval_expression_list([bin(col("a")).alias("result")])
    values = result.get_column_by_name("result").to_pylist()
    assert values == ["0", "1" + "0" * 63, "1" * 64]


def test_bin_float_input() -> None:
    from daft.functions import bin

    table = MicroPartition.from_pydict({"a": [3.5, 2.1]})
    with pytest.raises(ValueError, match="Expected input to bin to be integer"):
        table.eval_expression_list([bin(col("a"))])


def test_bin_bad_input() -> None:
    from daft.functions import bin

    table = MicroPartition.from_pydict({"a": ["a", "b", "c"]})
    with pytest.raises(ValueError, match="Expected input to bin to be integer"):
        table.eval_expression_list([bin(col("a"))])


@pytest.mark.parametrize(
    "value, from_base, to_base, expected",
    [
        pytest.param("3", 10, 2, "11", id="dec_to_bin"),
        pytest.param("FF", 16, 10, "255", id="hex_to_dec"),
        pytest.param("ff", 16, 10, "255", id="case_insensitive"),
        pytest.param("ZZ", 36, 10, "1295", id="max_base"),
        pytest.param("-15", 10, 16, "FFFFFFFFFFFFFFF1", id="neg_input_unsigned_out"),
        pytest.param("-15", 10, -16, "-F", id="neg_input_signed_out"),
        pytest.param("-1", 10, 16, "FFFFFFFFFFFFFFFF", id="neg_one_twos_complement"),
        pytest.param("11abc", 10, 16, "B", id="invalid_suffix_truncates"),
        pytest.param("ZZ", 16, 10, "0", id="all_invalid_chars"),
        pytest.param("Z", 10, 16, "0", id="leading_invalid"),
        pytest.param("-", 10, 16, "0", id="lone_minus_unsigned"),
        pytest.param("-", 10, -16, "-0", id="lone_minus_signed"),
        pytest.param("", 10, 16, None, id="empty"),
        pytest.param("  FF  ", 16, 10, "255", id="whitespace_trim"),
        pytest.param("FFFFFFFFFFFFFFFF", 16, 10, "18446744073709551615", id="u64_max"),
        pytest.param("18446744073709551616", 10, 16, None, id="u64_overflow"),
        pytest.param("-9223372036854775808", 10, 16, "8000000000000000", id="i64_min_two_complement"),
        pytest.param("-9223372036854775809", 10, 16, None, id="below_i64_min_returns_null"),
        pytest.param("-9223372036854775808", 10, -16, "-8000000000000000", id="i64_min_signed"),
        pytest.param(None, 16, 10, None, id="null_input"),
    ],
)
def test_conv(value: str | None, from_base: int, to_base: int, expected: str | None) -> None:
    from daft.functions import conv

    table = MicroPartition.from_pydict({"a": pa.array([value], type=pa.string())})
    result = table.eval_expression_list([conv(col("a"), from_base, to_base).alias("r")])
    assert result.get_column_by_name("r").to_pylist() == [expected]


def test_conv_int_input() -> None:
    from daft.functions import conv

    table = MicroPartition.from_pydict({"a": pa.array([255, 7, 0, None], type=pa.int64())})
    result = table.eval_expression_list([conv(col("a"), 10, 16).alias("r")])
    assert result.get_column_by_name("r").to_pylist() == ["FF", "7", "0", None]


def test_conv_unsigned_int_input() -> None:
    from daft.functions import conv

    table = MicroPartition.from_pydict({"a": pa.array([0, 2**63, 2**64 - 1], type=pa.uint64())})
    result = table.eval_expression_list([conv(col("a"), 10, 16).alias("r")])
    assert result.get_column_by_name("r").to_pylist() == ["0", "8000000000000000", "FFFFFFFFFFFFFFFF"]


def test_conv_negative_int_input() -> None:
    from daft.functions import conv

    table = MicroPartition.from_pydict({"a": pa.array([-1, -255, -(2**63)], type=pa.int64())})
    result = table.eval_expression_list([conv(col("a"), 10, 16).alias("r")])
    assert result.get_column_by_name("r").to_pylist() == [
        "FFFFFFFFFFFFFFFF",
        "FFFFFFFFFFFFFF01",
        "8000000000000000",
    ]


@pytest.mark.parametrize(
    "from_base, to_base",
    [
        pytest.param(1, 10, id="from_too_small"),
        pytest.param(37, 10, id="from_too_large"),
        pytest.param(10, 0, id="to_zero"),
        pytest.param(10, 1, id="to_too_small"),
        pytest.param(10, 37, id="to_too_large"),
        pytest.param(10, -37, id="to_too_negative"),
        pytest.param(-10, 16, id="from_negative"),
    ],
)
def test_conv_invalid_base(from_base: int, to_base: int) -> None:
    from daft.functions import conv

    table = MicroPartition.from_pydict({"a": ["FF"]})
    result = table.eval_expression_list([conv(col("a"), from_base, to_base).alias("r")])
    assert result.get_column_by_name("r").to_pylist() == [None]


def test_conv_bad_dtype() -> None:
    from daft.functions import conv

    table = MicroPartition.from_pydict({"a": [1.0, 2.5]})
    with pytest.raises(ValueError, match="Expected input to conv to be string or integer"):
        table.eval_expression_list([conv(col("a"), 10, 16)])


def test_hypot() -> None:
    from daft.functions import hypot

    table = MicroPartition.from_pydict({"a": [3.0, 5.0, 0.0, None], "b": [4.0, 12.0, 0.0, None]})
    result = table.eval_expression_list([hypot(col("a"), col("b")).alias("result")])
    values = result.get_column_by_name("result").to_pylist()
    expected = [5.0, 13.0, 0.0, None]
    assert values == expected


def test_hypot_integers() -> None:
    from daft.functions import hypot

    table = MicroPartition.from_pydict({"a": [3, 5], "b": [4, 12]})
    result = table.eval_expression_list([hypot(col("a"), col("b")).alias("result")])
    values = result.get_column_by_name("result").to_pylist()
    expected = [5.0, 13.0]
    assert values == expected


def test_hypot_bad_input() -> None:
    from daft.functions import hypot

    table = MicroPartition.from_pydict({"a": ["a", "b"], "b": ["c", "d"]})
    with pytest.raises(ValueError, match="Expected inputs to hypot to be numeric"):
        table.eval_expression_list([hypot(col("a"), col("b"))])


def test_pmod() -> None:
    from daft.functions import pmod

    table = MicroPartition.from_pydict({"a": [7, 10, 0, 5], "b": [3, 4, 5, 5]})
    result = table.eval_expression_list([pmod(col("a"), col("b")).alias("result")])
    values = result.get_column_by_name("result").to_pylist()
    assert values == [1, 2, 0, 0]


def test_pmod_negative() -> None:
    from daft.functions import pmod

    table = MicroPartition.from_pydict(
        {"a": pa.array([-7, 7, -7, -8], type=pa.int64()), "b": pa.array([3, -3, -3, -3], type=pa.int64())}
    )
    result = table.eval_expression_list([pmod(col("a"), col("b")).alias("result")])
    values = result.get_column_by_name("result").to_pylist()
    assert values == [2, 1, -1, -2]


def test_pmod_float() -> None:
    from daft.functions import pmod

    table = MicroPartition.from_pydict({"a": [7.5, -7.0, 1.0, -1.0], "b": [2.5, 3.0, -3.0, -3.0]})
    result = table.eval_expression_list([pmod(col("a"), col("b")).alias("result")])
    values = result.get_column_by_name("result").to_pylist()
    assert values == [0.0, 2.0, 1.0, -1.0]


def test_pmod_zero_divisor() -> None:
    from daft.functions import pmod

    table = MicroPartition.from_pydict({"a": [7, 0, -3, 5], "b": [0, 0, 0, 2]})
    result = table.eval_expression_list([pmod(col("a"), col("b")).alias("result")])
    values = result.get_column_by_name("result").to_pylist()
    assert values == [None, None, None, 1]


def test_pmod_int_boundary() -> None:
    from daft.functions import pmod

    # i64::MIN exercises the wrapping path: i64::MIN % -1 would panic without
    # wrapping_rem; pmod(i64::MIN, -1) is mathematically 0.
    table = MicroPartition.from_pydict(
        {"a": pa.array([-(2**63), -(2**63)], type=pa.int64()), "b": pa.array([-1, 3], type=pa.int64())}
    )
    result = table.eval_expression_list([pmod(col("a"), col("b")).alias("result")])
    values = result.get_column_by_name("result").to_pylist()
    assert values == [0, 1]


@pytest.mark.parametrize(
    "a, b, expected",
    [
        pytest.param(float("nan"), 3.0, float("nan"), id="nan_dividend"),
        pytest.param(1.0, float("nan"), float("nan"), id="nan_divisor"),
        pytest.param(float("inf"), 3.0, float("nan"), id="inf_dividend"),
        pytest.param(1.0, float("inf"), 1.0, id="inf_divisor_positive_a"),
        pytest.param(-1.0, float("inf"), float("nan"), id="inf_divisor_negative_a"),
    ],
)
def test_pmod_float_nan_inf(a: float, b: float, expected: float) -> None:
    from daft.functions import pmod

    table = MicroPartition.from_pydict({"a": [a], "b": [b]})
    result = table.eval_expression_list([pmod(col("a"), col("b")).alias("result")])
    [value] = result.get_column_by_name("result").to_pylist()
    if math.isnan(expected):
        assert math.isnan(value)
    else:
        assert value == expected


@pytest.mark.parametrize(
    "a_arrow, b_arrow",
    [
        pytest.param(pa.int8(), pa.int8(), id="i8/i8"),
        pytest.param(pa.int32(), pa.int64(), id="i32/i64"),
        pytest.param(pa.int8(), pa.int32(), id="i8/i32"),
        pytest.param(pa.uint16(), pa.uint32(), id="u16/u32"),
    ],
)
def test_pmod_int_cross_width(a_arrow: pa.DataType, b_arrow: pa.DataType) -> None:
    from daft.functions import pmod

    table = MicroPartition.from_pydict({"a": pa.array([7, 5, 0], type=a_arrow), "b": pa.array([3, 4, 2], type=b_arrow)})
    result = table.eval_expression_list([pmod(col("a"), col("b")).alias("result")])
    assert result.get_column_by_name("result").to_pylist() == [1, 1, 0]


def test_pmod_bad_input() -> None:
    from daft.functions import pmod

    table = MicroPartition.from_pydict({"a": ["x", "y"], "b": [1, 2]})
    with pytest.raises(ValueError, match="Expected inputs to pmod to be numeric"):
        table.eval_expression_list([pmod(col("a"), col("b"))])


def test_width_bucket_basic() -> None:
    table = MicroPartition.from_pydict(
        {
            "v": [5.3, -2.1, 8.1, -0.9],
            "mn": [0.2, 1.3, 0.0, 5.2],
            "mx": [10.6, 3.4, 5.7, 0.5],
            "nb": [5, 3, 4, 2],
        }
    )
    result = table.eval_expression_list([width_bucket(col("v"), col("mn"), col("mx"), col("nb")).alias("bucket")])
    assert result.get_column_by_name("bucket").to_pylist() == [3, 0, 5, 3]


def test_width_bucket_descending() -> None:
    """Min > max flips orientation; v > min underflows, v <= max overflows."""
    table = MicroPartition.from_pydict({"v": [11.0, 5.0, 0.0, -1.0]})
    result = table.eval_expression_list([width_bucket(col("v"), lit(10.0), lit(0.0), lit(5)).alias("bucket")])
    # min=10, max=0, nb=5: v=11 underflows -> 0; v=5 mid -> floor(5*(10-5)/10)+1=3;
    # v=0 == max -> nb+1=6; v=-1 < max -> nb+1=6.
    assert result.get_column_by_name("bucket").to_pylist() == [0, 3, 6, 6]


def test_width_bucket_integer_inputs() -> None:
    table = MicroPartition.from_pydict({"v": pa.array([3, 5, 7], type=pa.int32())})
    result = table.eval_expression_list([width_bucket(col("v"), lit(0), lit(10), lit(5)).alias("bucket")])
    assert result.get_column_by_name("bucket").to_pylist() == [2, 3, 4]


def test_width_bucket_below_and_above_range() -> None:
    table = MicroPartition.from_pydict({"v": [-0.1, 0.0, 5.0, 10.0, 10.1]})
    result = table.eval_expression_list([width_bucket(col("v"), lit(0.0), lit(10.0), lit(5)).alias("bucket")])
    # v=0 is at min so bucket 1; v=10 hits the v>=max branch so nb+1=6.
    assert result.get_column_by_name("bucket").to_pylist() == [0, 1, 3, 6, 6]


def test_width_bucket_null_returning() -> None:
    """Each invalid row exercises a distinct NULL-returning guard in compute_bucket."""
    i64_max = 2**63 - 1
    nan, inf = math.nan, math.inf
    table = MicroPartition.from_pydict(
        {
            #     nb<=0  nb==i64::MAX  mn==mx  v=NaN  mn=NaN  mn=Inf  mx=NaN  mx=Inf  valid
            "v": [1.0, 1.0, 1.0, nan, 1.0, 1.0, 1.0, 1.0, 1.0],
            "mn": [0.0, 0.0, 5.0, 0.0, nan, inf, 0.0, 0.0, 0.0],
            "mx": [10.0, 10.0, 5.0, 10.0, 10.0, 10.0, nan, inf, 10.0],
            "nb": [0, i64_max, 4, 4, 4, 4, 4, 4, 4],
        }
    )
    result = table.eval_expression_list([width_bucket(col("v"), col("mn"), col("mx"), col("nb")).alias("bucket")])
    values = result.get_column_by_name("bucket").to_pylist()
    assert values[:-1] == [None] * 8
    assert values[-1] is not None


def test_width_bucket_null_propagation() -> None:
    table = MicroPartition.from_pydict({"v": [1.0, None, 5.0]})
    result = table.eval_expression_list([width_bucket(col("v"), lit(0.0), lit(10.0), lit(5)).alias("bucket")])
    assert result.get_column_by_name("bucket").to_pylist() == [1, None, 3]


def test_width_bucket_scalar_broadcast() -> None:
    table = MicroPartition.from_pydict({"v": [-1.0, 0.0, 2.0, 4.0, 6.0, 8.0, 10.0, 11.0]})
    result = table.eval_expression_list([width_bucket(col("v"), lit(0.0), lit(10.0), lit(5)).alias("bucket")])
    assert result.get_column_by_name("bucket").to_pylist() == [0, 1, 2, 3, 4, 5, 6, 6]


def test_width_bucket_fractional_num_bucket_truncates() -> None:
    """Daft promotes any numeric num_bucket to Int64 by truncation."""
    table = MicroPartition.from_pydict({"v": [3.0]})
    result = table.eval_expression_list([width_bucket(col("v"), lit(0.0), lit(10.0), lit(5.7)).alias("bucket")])
    # 5.7 truncates to 5; floor(5 * 3 / 10) + 1 = 2.
    assert result.get_column_by_name("bucket").to_pylist() == [2]


def test_width_bucket_bad_input() -> None:
    table = MicroPartition.from_pydict({"v": ["a", "b"]})
    with pytest.raises(ValueError, match="Expected `value` of width_bucket to be numeric"):
        table.eval_expression_list([width_bucket(col("v"), lit(0.0), lit(10.0), lit(5))])


def test_width_bucket_huge_num_bucket() -> None:
    """Nb just below i64::MAX must saturate, not wrap to i64::MIN."""
    nb = 2**63 - 2  # one below the i64::MAX None-guard; (nb as f64) rounds up to 2^63
    table = MicroPartition.from_pydict({"v": [-1.0, 5.0, 11.0]})
    result = table.eval_expression_list([width_bucket(col("v"), lit(0.0), lit(10.0), lit(nb)).alias("bucket")])
    values = result.get_column_by_name("bucket").to_pylist()
    assert values[0] == 0  # below range
    assert values[1] == 2**62 + 1  # 2^63 * 0.5 truncated to i64, then saturating_add(1)
    assert values[2] == 2**63 - 1  # at/above range: nb.saturating_add(1) -> i64::MAX