# 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 math
import warnings
from datetime import date, datetime, time, timezone

import numpy as np
import pyarrow as pa
import pytest
from datafusion import SessionContext, column, literal
from datafusion import functions as f
from datafusion.expr import GroupingSet

np.seterr(invalid="ignore")

DEFAULT_TZ = timezone.utc


@pytest.fixture
def df():
    ctx = SessionContext()
    # create a RecordBatch and a new DataFrame from it
    batch = pa.RecordBatch.from_arrays(
        [
            pa.array(["Hello", "World", "!"], type=pa.string_view()),
            pa.array([4, 5, 6]),
            pa.array(["hello ", " world ", " !"], type=pa.string_view()),
            pa.array(
                [
                    datetime(2022, 12, 31, tzinfo=DEFAULT_TZ),
                    datetime(2027, 6, 26, tzinfo=DEFAULT_TZ),
                    datetime(2020, 7, 2, tzinfo=DEFAULT_TZ),
                ]
            ),
            pa.array([False, True, True]),
        ],
        names=["a", "b", "c", "d", "e"],
    )
    return ctx.create_dataframe([[batch]])


def test_named_struct(df):
    df = df.with_column(
        "d",
        f.named_struct(
            [
                ("a", column("a")),
                ("b", column("b")),
                ("c", column("c")),
            ]
        ),
    )

    expected = """DataFrame()
+-------+---+---------+------------------------------+-------+
| a     | b | c       | d                            | e     |
+-------+---+---------+------------------------------+-------+
| Hello | 4 | hello   | {a: Hello, b: 4, c: hello }  | false |
| World | 5 |  world  | {a: World, b: 5, c:  world } | true  |
| !     | 6 |  !      | {a: !, b: 6, c:  !}          | true  |
+-------+---+---------+------------------------------+-------+
""".strip()
    assert str(df) == expected


def test_literal(df):
    df = df.select(
        literal(1),
        literal("1"),
        literal("OK"),
        literal(3.14),
        literal(value=True),
        literal(b"hello world"),
    )
    result = df.collect()
    assert len(result) == 1
    result = result[0]
    assert result.column(0) == pa.array([1] * 3)
    assert result.column(1) == pa.array(["1"] * 3, type=pa.string_view())
    assert result.column(2) == pa.array(["OK"] * 3, type=pa.string_view())
    assert result.column(3) == pa.array([3.14] * 3)
    assert result.column(4) == pa.array([True] * 3)
    assert result.column(5) == pa.array([b"hello world"] * 3)


def test_lit_arith(df):
    """Test literals with arithmetic operations"""
    df = df.select(
        literal(1) + column("b"), f.concat(column("a").cast(pa.string()), literal("!"))
    )
    result = df.collect()
    assert len(result) == 1
    result = result[0]

    assert result.column(0) == pa.array([5, 6, 7])
    assert result.column(1) == pa.array(
        ["Hello!", "World!", "!!"], type=pa.string_view()
    )


def test_math_functions():
    ctx = SessionContext()
    # create a RecordBatch and a new DataFrame from it
    batch = pa.RecordBatch.from_arrays(
        [pa.array([0.1, -0.7, 0.55]), pa.array([float("nan"), 0, 2.0])],
        names=["value", "na_value"],
    )
    df = ctx.create_dataframe([[batch]])

    values = np.array([0.1, -0.7, 0.55])
    na_values = np.array([np.nan, 0, 2.0])
    col_v = column("value")
    col_nav = column("na_value")
    df = df.select(
        f.abs(col_v),
        f.sin(col_v),
        f.cos(col_v),
        f.tan(col_v),
        f.asin(col_v),
        f.acos(col_v),
        f.exp(col_v),
        f.ln(col_v + literal(pa.scalar(1))),
        f.log2(col_v + literal(pa.scalar(1))),
        f.log10(col_v + literal(pa.scalar(1))),
        f.random(),
        f.atan(col_v),
        f.atan2(col_v, literal(pa.scalar(1.1))),
        f.ceil(col_v),
        f.floor(col_v),
        f.power(col_v, literal(pa.scalar(3))),
        f.pow(col_v, literal(pa.scalar(4))),
        f.round(col_v),
        f.round(col_v, literal(pa.scalar(3))),
        f.sqrt(col_v),
        f.signum(col_v),
        f.trunc(col_v),
        f.asinh(col_v),
        f.acosh(col_v),
        f.atanh(col_v),
        f.cbrt(col_v),
        f.cosh(col_v),
        f.degrees(col_v),
        f.gcd(literal(9), literal(3)),
        f.lcm(literal(6), literal(4)),
        f.nanvl(col_nav, literal(5)),
        f.pi(),
        f.radians(col_v),
        f.sinh(col_v),
        f.tanh(col_v),
        f.factorial(literal(6)),
        f.isnan(col_nav),
        f.iszero(col_nav),
        f.log(literal(3), col_v + literal(pa.scalar(1))),
    )
    batches = df.collect()
    assert len(batches) == 1
    result = batches[0]

    np.testing.assert_array_almost_equal(result.column(0), np.abs(values))
    np.testing.assert_array_almost_equal(result.column(1), np.sin(values))
    np.testing.assert_array_almost_equal(result.column(2), np.cos(values))
    np.testing.assert_array_almost_equal(result.column(3), np.tan(values))
    np.testing.assert_array_almost_equal(result.column(4), np.arcsin(values))
    np.testing.assert_array_almost_equal(result.column(5), np.arccos(values))
    np.testing.assert_array_almost_equal(result.column(6), np.exp(values))
    np.testing.assert_array_almost_equal(result.column(7), np.log(values + 1.0))
    np.testing.assert_array_almost_equal(result.column(8), np.log2(values + 1.0))
    np.testing.assert_array_almost_equal(result.column(9), np.log10(values + 1.0))
    np.testing.assert_array_less(result.column(10), np.ones_like(values))
    np.testing.assert_array_almost_equal(result.column(11), np.arctan(values))
    np.testing.assert_array_almost_equal(result.column(12), np.arctan2(values, 1.1))
    np.testing.assert_array_almost_equal(result.column(13), np.ceil(values))
    np.testing.assert_array_almost_equal(result.column(14), np.floor(values))
    np.testing.assert_array_almost_equal(result.column(15), np.power(values, 3))
    np.testing.assert_array_almost_equal(result.column(16), np.power(values, 4))
    np.testing.assert_array_almost_equal(result.column(17), np.round(values))
    np.testing.assert_array_almost_equal(result.column(18), np.round(values, 3))
    np.testing.assert_array_almost_equal(result.column(19), np.sqrt(values))
    np.testing.assert_array_almost_equal(result.column(20), np.sign(values))
    np.testing.assert_array_almost_equal(result.column(21), np.trunc(values))
    np.testing.assert_array_almost_equal(result.column(22), np.arcsinh(values))
    np.testing.assert_array_almost_equal(result.column(23), np.arccosh(values))
    np.testing.assert_array_almost_equal(result.column(24), np.arctanh(values))
    np.testing.assert_array_almost_equal(result.column(25), np.cbrt(values))
    np.testing.assert_array_almost_equal(result.column(26), np.cosh(values))
    np.testing.assert_array_almost_equal(result.column(27), np.degrees(values))
    np.testing.assert_array_almost_equal(result.column(28), np.gcd(9, 3))
    np.testing.assert_array_almost_equal(result.column(29), np.lcm(6, 4))
    np.testing.assert_array_almost_equal(
        result.column(30), np.where(np.isnan(na_values), 5, na_values)
    )
    np.testing.assert_array_almost_equal(result.column(31), np.pi)
    np.testing.assert_array_almost_equal(result.column(32), np.radians(values))
    np.testing.assert_array_almost_equal(result.column(33), np.sinh(values))
    np.testing.assert_array_almost_equal(result.column(34), np.tanh(values))
    np.testing.assert_array_almost_equal(result.column(35), math.factorial(6))
    np.testing.assert_array_almost_equal(result.column(36), np.isnan(na_values))
    np.testing.assert_array_almost_equal(result.column(37), na_values == 0)
    np.testing.assert_array_almost_equal(
        result.column(38), np.emath.logn(3, values + 1.0)
    )


def py_indexof(arr, v):
    try:
        return arr.index(v) + 1
    except ValueError:
        return np.nan


def py_arr_remove(arr, v, n=None):
    new_arr = arr[:]
    found = 0
    try:
        while found != n:
            new_arr.remove(v)
            found += 1
    except ValueError:
        pass

    return new_arr


def py_arr_replace(arr, from_, to, n=None):
    new_arr = arr[:]
    found = 0
    try:
        while found != n:
            idx = new_arr.index(from_)
            new_arr[idx] = to
            found += 1
    except ValueError:
        pass

    return new_arr


def py_arr_resize(arr, size, value):
    arr = np.asarray(arr)
    return np.pad(
        arr,
        [(0, size - arr.shape[0])],
        "constant",
        constant_values=value,
    )


def py_flatten(arr):
    result = []
    for elem in arr:
        if isinstance(elem, list):
            result.extend(py_flatten(elem))
        else:
            result.append(elem)
    return result


@pytest.mark.parametrize(
    ("stmt", "py_expr"),
    [
        (
            lambda col: f.array_append(col, literal(99.0)),
            lambda data: [np.append(arr, 99.0) for arr in data],
        ),
        (
            lambda col: f.array_push_back(col, literal(99.0)),
            lambda data: [np.append(arr, 99.0) for arr in data],
        ),
        (
            lambda col: f.list_append(col, literal(99.0)),
            lambda data: [np.append(arr, 99.0) for arr in data],
        ),
        (
            lambda col: f.list_push_back(col, literal(99.0)),
            lambda data: [np.append(arr, 99.0) for arr in data],
        ),
        (
            lambda col: f.array_concat(col, col),
            lambda data: [np.concatenate([arr, arr]) for arr in data],
        ),
        (
            lambda col: f.array_cat(col, col),
            lambda data: [np.concatenate([arr, arr]) for arr in data],
        ),
        (
            lambda col: f.list_cat(col, col),
            lambda data: [np.concatenate([arr, arr]) for arr in data],
        ),
        (
            lambda col: f.list_concat(col, col),
            lambda data: [np.concatenate([arr, arr]) for arr in data],
        ),
        (
            f.array_dims,
            lambda data: [[len(r)] for r in data],
        ),
        (
            lambda col: f.array_sort(f.array_distinct(col)),
            lambda data: [sorted(set(r)) for r in data],
        ),
        (
            lambda col: f.list_sort(f.list_distinct(col)),
            lambda data: [sorted(set(r)) for r in data],
        ),
        (
            f.list_dims,
            lambda data: [[len(r)] for r in data],
        ),
        (
            lambda col: f.array_element(col, literal(1)),
            lambda data: [r[0] for r in data],
        ),
        (
            f.array_empty,
            lambda data: [len(r) == 0 for r in data],
        ),
        (
            f.empty,
            lambda data: [len(r) == 0 for r in data],
        ),
        (
            f.list_empty,
            lambda data: [len(r) == 0 for r in data],
        ),
        (
            lambda col: f.array_extract(col, literal(1)),
            lambda data: [r[0] for r in data],
        ),
        (
            lambda col: f.list_element(col, literal(1)),
            lambda data: [r[0] for r in data],
        ),
        (
            lambda col: f.list_extract(col, literal(1)),
            lambda data: [r[0] for r in data],
        ),
        (
            f.array_length,
            lambda data: [len(r) for r in data],
        ),
        (
            f.list_length,
            lambda data: [len(r) for r in data],
        ),
        (
            lambda col: f.array_has(col, literal(1.0)),
            lambda data: [1.0 in r for r in data],
        ),
        (
            lambda col: f.list_has(col, literal(1.0)),
            lambda data: [1.0 in r for r in data],
        ),
        (
            lambda col: f.array_contains(col, literal(1.0)),
            lambda data: [1.0 in r for r in data],
        ),
        (
            lambda col: f.list_contains(col, literal(1.0)),
            lambda data: [1.0 in r for r in data],
        ),
        (
            lambda col: f.array_has_all(
                col, f.make_array(*[literal(v) for v in [1.0, 3.0, 5.0]])
            ),
            lambda data: [np.all([v in r for v in [1.0, 3.0, 5.0]]) for r in data],
        ),
        (
            lambda col: f.list_has_all(
                col, f.make_array(*[literal(v) for v in [1.0, 3.0, 5.0]])
            ),
            lambda data: [np.all([v in r for v in [1.0, 3.0, 5.0]]) for r in data],
        ),
        (
            lambda col: f.array_has_any(
                col, f.make_array(*[literal(v) for v in [1.0, 3.0, 5.0]])
            ),
            lambda data: [np.any([v in r for v in [1.0, 3.0, 5.0]]) for r in data],
        ),
        (
            lambda col: f.list_has_any(
                col, f.make_array(*[literal(v) for v in [1.0, 3.0, 5.0]])
            ),
            lambda data: [np.any([v in r for v in [1.0, 3.0, 5.0]]) for r in data],
        ),
        (
            lambda col: f.arrays_overlap(
                col, f.make_array(*[literal(v) for v in [1.0, 3.0, 5.0]])
            ),
            lambda data: [np.any([v in r for v in [1.0, 3.0, 5.0]]) for r in data],
        ),
        (
            lambda col: f.list_overlap(
                col, f.make_array(*[literal(v) for v in [1.0, 3.0, 5.0]])
            ),
            lambda data: [np.any([v in r for v in [1.0, 3.0, 5.0]]) for r in data],
        ),
        (
            lambda col: f.array_position(col, literal(1.0)),
            lambda data: [py_indexof(r, 1.0) for r in data],
        ),
        (
            lambda col: f.array_indexof(col, literal(1.0)),
            lambda data: [py_indexof(r, 1.0) for r in data],
        ),
        (
            lambda col: f.list_position(col, literal(1.0)),
            lambda data: [py_indexof(r, 1.0) for r in data],
        ),
        (
            lambda col: f.list_indexof(col, literal(1.0)),
            lambda data: [py_indexof(r, 1.0) for r in data],
        ),
        (
            lambda col: f.array_positions(col, literal(1.0)),
            lambda data: [[i + 1 for i, _v in enumerate(r) if _v == 1.0] for r in data],
        ),
        (
            lambda col: f.list_positions(col, literal(1.0)),
            lambda data: [[i + 1 for i, _v in enumerate(r) if _v == 1.0] for r in data],
        ),
        (
            f.array_ndims,
            lambda data: [np.array(r).ndim for r in data],
        ),
        (
            f.list_ndims,
            lambda data: [np.array(r).ndim for r in data],
        ),
        (
            lambda col: f.array_prepend(literal(99.0), col),
            lambda data: [np.insert(arr, 0, 99.0) for arr in data],
        ),
        (
            lambda col: f.array_push_front(literal(99.0), col),
            lambda data: [np.insert(arr, 0, 99.0) for arr in data],
        ),
        (
            lambda col: f.list_prepend(literal(99.0), col),
            lambda data: [np.insert(arr, 0, 99.0) for arr in data],
        ),
        (
            lambda col: f.list_push_front(literal(99.0), col),
            lambda data: [np.insert(arr, 0, 99.0) for arr in data],
        ),
        (
            f.array_pop_back,
            lambda data: [arr[:-1] for arr in data],
        ),
        (
            f.list_pop_back,
            lambda data: [arr[:-1] for arr in data],
        ),
        (
            f.array_pop_front,
            lambda data: [arr[1:] for arr in data],
        ),
        (
            f.list_pop_front,
            lambda data: [arr[1:] for arr in data],
        ),
        (
            lambda col: f.array_remove(col, literal(3.0)),
            lambda data: [py_arr_remove(arr, 3.0, 1) for arr in data],
        ),
        (
            lambda col: f.list_remove(col, literal(3.0)),
            lambda data: [py_arr_remove(arr, 3.0, 1) for arr in data],
        ),
        (
            lambda col: f.array_remove_n(col, literal(3.0), literal(2)),
            lambda data: [py_arr_remove(arr, 3.0, 2) for arr in data],
        ),
        (
            lambda col: f.list_remove_n(col, literal(3.0), literal(2)),
            lambda data: [py_arr_remove(arr, 3.0, 2) for arr in data],
        ),
        (
            lambda col: f.array_remove_all(col, literal(3.0)),
            lambda data: [py_arr_remove(arr, 3.0) for arr in data],
        ),
        (
            lambda col: f.list_remove_all(col, literal(3.0)),
            lambda data: [py_arr_remove(arr, 3.0) for arr in data],
        ),
        (
            lambda col: f.array_repeat(col, literal(2)),
            lambda data: [[arr] * 2 for arr in data],
        ),
        (
            lambda col: f.list_repeat(col, literal(2)),
            lambda data: [[arr] * 2 for arr in data],
        ),
        (
            lambda col: f.array_replace(col, literal(3.0), literal(4.0)),
            lambda data: [py_arr_replace(arr, 3.0, 4.0, 1) for arr in data],
        ),
        (
            lambda col: f.list_replace(col, literal(3.0), literal(4.0)),
            lambda data: [py_arr_replace(arr, 3.0, 4.0, 1) for arr in data],
        ),
        (
            lambda col: f.array_replace_n(col, literal(3.0), literal(4.0), literal(1)),
            lambda data: [py_arr_replace(arr, 3.0, 4.0, 1) for arr in data],
        ),
        (
            lambda col: f.list_replace_n(col, literal(3.0), literal(4.0), literal(2)),
            lambda data: [py_arr_replace(arr, 3.0, 4.0, 2) for arr in data],
        ),
        (
            lambda col: f.array_replace_all(col, literal(3.0), literal(4.0)),
            lambda data: [py_arr_replace(arr, 3.0, 4.0) for arr in data],
        ),
        (
            lambda col: f.list_replace_all(col, literal(3.0), literal(4.0)),
            lambda data: [py_arr_replace(arr, 3.0, 4.0) for arr in data],
        ),
        (
            lambda col: f.array_sort(col, descending=True, null_first=True),
            lambda data: [np.sort(arr)[::-1] for arr in data],
        ),
        (
            lambda col: f.list_sort(col, descending=False, null_first=False),
            lambda data: [np.sort(arr) for arr in data],
        ),
        (
            lambda col: f.array_slice(col, literal(2), literal(4)),
            lambda data: [arr[1:4] for arr in data],
        ),
        pytest.param(
            lambda col: f.list_slice(col, literal(-1), literal(2)),
            lambda data: [arr[-1:2] for arr in data],
        ),
        (
            lambda col: col[:3],
            lambda data: [arr[:3] for arr in data],
        ),
        (
            lambda col: col[1:3],
            lambda data: [arr[1:3] for arr in data],
        ),
        (
            lambda col: col[1:4:2],
            lambda data: [arr[1:4:2] for arr in data],
        ),
        (
            lambda col: col[literal(1) : literal(4)],
            lambda data: [arr[1:4] for arr in data],
        ),
        (
            lambda col: col[column("indices") : column("indices") + literal(2)],
            lambda data: [[2.0, 3.0], [], [6.0]],
        ),
        (
            lambda col: col[literal(1) : literal(4) : literal(2)],
            lambda data: [arr[1:4:2] for arr in data],
        ),
        (
            lambda col: f.array_sort(f.array_intersect(col, literal([3.0, 4.0]))),
            lambda data: [np.intersect1d(arr, [3.0, 4.0]) for arr in data],
        ),
        (
            lambda col: f.list_sort(f.list_intersect(col, literal([3.0, 4.0]))),
            lambda data: [np.intersect1d(arr, [3.0, 4.0]) for arr in data],
        ),
        (
            lambda col: f.array_sort(f.array_union(col, literal([12.0, 999.0]))),
            lambda data: [np.union1d(arr, [12.0, 999.0]) for arr in data],
        ),
        (
            lambda col: f.list_sort(f.list_union(col, literal([12.0, 999.0]))),
            lambda data: [np.union1d(arr, [12.0, 999.0]) for arr in data],
        ),
        (
            lambda col: f.array_except(col, literal([3.0])),
            lambda data: [np.setdiff1d(arr, [3.0]) for arr in data],
        ),
        (
            lambda col: f.list_except(col, literal([3.0])),
            lambda data: [np.setdiff1d(arr, [3.0]) for arr in data],
        ),
        (
            lambda col: f.array_resize(col, literal(10), literal(0.0)),
            lambda data: [py_arr_resize(arr, 10, 0.0) for arr in data],
        ),
        (
            lambda col: f.list_resize(col, literal(10), literal(0.0)),
            lambda data: [py_arr_resize(arr, 10, 0.0) for arr in data],
        ),
        (
            lambda col: f.range(literal(1), literal(5), literal(2)),
            lambda data: [np.arange(1, 5, 2)],
        ),
    ],
)
def test_array_functions(stmt, py_expr):
    data = [[1.0, 2.0, 3.0, 3.0], [4.0, 5.0, 3.0], [6.0]]
    indices = [1, 3, 0]
    ctx = SessionContext()
    batch = pa.RecordBatch.from_arrays(
        [np.array(data, dtype=object), indices], names=["arr", "indices"]
    )
    df = ctx.create_dataframe([[batch]])

    col = column("arr")
    query_result = df.select(stmt(col)).collect()[0].column(0)
    for a, b in zip(query_result, py_expr(data), strict=False):
        np.testing.assert_array_almost_equal(
            np.array(a.as_py(), dtype=float), np.array(b, dtype=float)
        )


def test_array_function_flatten():
    data = [[1.0, 2.0, 3.0, 3.0], [4.0, 5.0, 3.0], [6.0]]
    ctx = SessionContext()
    batch = pa.RecordBatch.from_arrays([np.array(data, dtype=object)], names=["arr"])
    df = ctx.create_dataframe([[batch]])

    stmt = f.flatten(literal(data))
    py_expr = [py_flatten(data)]
    query_result = df.select(stmt).collect()[0].column(0)
    for a, b in zip(query_result, py_expr, strict=False):
        np.testing.assert_array_almost_equal(
            np.array(a.as_py(), dtype=float), np.array(b, dtype=float)
        )


def test_array_function_cardinality():
    data = [[1, 2, 3], [4, 4, 5, 6]]
    ctx = SessionContext()
    batch = pa.RecordBatch.from_arrays([np.array(data, dtype=object)], names=["arr"])
    df = ctx.create_dataframe([[batch]])

    stmt = f.cardinality(column("arr"))
    py_expr = [len(arr) for arr in data]  # Expected lengths: [3, 3]
    # assert py_expr lengths

    query_result = df.select(stmt).collect()[0].column(0)

    for a, b in zip(query_result, py_expr, strict=False):
        np.testing.assert_array_equal(
            np.array([a.as_py()], dtype=int), np.array([b], dtype=int)
        )


@pytest.mark.parametrize("make_func", [f.make_array, f.make_list])
def test_make_array_functions(make_func):
    ctx = SessionContext()
    batch = pa.RecordBatch.from_arrays(
        [
            pa.array(["Hello", "World", "!"], type=pa.string()),
            pa.array([4, 5, 6]),
            pa.array(["hello ", " world ", " !"], type=pa.string()),
        ],
        names=["a", "b", "c"],
    )
    df = ctx.create_dataframe([[batch]])

    stmt = make_func(
        column("a").cast(pa.string()),
        column("b").cast(pa.string()),
        column("c").cast(pa.string()),
    )
    py_expr = [
        ["Hello", "4", "hello "],
        ["World", "5", " world "],
        ["!", "6", " !"],
    ]

    query_result = df.select(stmt).collect()[0].column(0)
    for a, b in zip(query_result, py_expr, strict=False):
        np.testing.assert_array_equal(
            np.array(a.as_py(), dtype=str), np.array(b, dtype=str)
        )


@pytest.mark.parametrize(
    ("stmt", "py_expr"),
    [
        (
            f.array_to_string(column("arr"), literal(",")),
            lambda data: [",".join([str(int(v)) for v in r]) for r in data],
        ),
        (
            f.array_join(column("arr"), literal(",")),
            lambda data: [",".join([str(int(v)) for v in r]) for r in data],
        ),
        (
            f.list_to_string(column("arr"), literal(",")),
            lambda data: [",".join([str(int(v)) for v in r]) for r in data],
        ),
        (
            f.list_join(column("arr"), literal(",")),
            lambda data: [",".join([str(int(v)) for v in r]) for r in data],
        ),
    ],
)
def test_array_function_obj_tests(stmt, py_expr):
    data = [[1.0, 2.0, 3.0, 3.0], [4.0, 5.0, 3.0], [6.0]]
    ctx = SessionContext()
    batch = pa.RecordBatch.from_arrays([np.array(data, dtype=object)], names=["arr"])
    df = ctx.create_dataframe([[batch]])
    query_result = np.array(df.select(stmt).collect()[0].column(0))
    for a, b in zip(query_result, py_expr(data), strict=False):
        assert a == b


@pytest.mark.parametrize(
    ("alias_fn", "primary_fn", "data"),
    [
        (f.list_compact, f.array_compact, {"a": [[1.0, None, 2.0, None, 3.0]]}),
        (f.list_normalize, f.array_normalize, {"a": [[3.0, 4.0]]}),
        (
            f.dot_product,
            f.inner_product,
            {"a": [[1.0, 2.0, 3.0]], "b": [[4.0, 5.0, 6.0]]},
        ),
    ],
)
def test_array_function_aliases(alias_fn, primary_fn, data):
    """Alias helpers should be exact aliases for their primary counterparts."""
    ctx = SessionContext()
    df = ctx.from_pydict(data)
    cols = [column(name) for name in data]
    alias_result = df.select(alias_fn(*cols).alias("r")).collect()
    primary_result = df.select(primary_fn(*cols).alias("r")).collect()
    assert (
        alias_result[0].column(0).to_pylist() == primary_result[0].column(0).to_pylist()
    )


@pytest.mark.parametrize("fn", [f.cosine_distance, f.inner_product, f.dot_product])
def test_array_distance_length_mismatch_raises(fn):
    """Length-mismatched inputs to vector distance fns should raise at execute."""
    ctx = SessionContext()
    df = ctx.from_pydict({"a": [[1.0, 2.0]], "b": [[1.0, 2.0, 3.0]]})
    with pytest.raises(Exception, match="same length"):
        df.select(fn(column("a"), column("b")).alias("r")).collect()


@pytest.mark.parametrize(
    ("args", "expected"),
    [
        pytest.param(
            ({"x": 1, "y": 2},),
            [("x", 1), ("y", 2)],
            id="dict",
        ),
        pytest.param(
            ({"x": literal(1), "y": literal(2)},),
            [("x", 1), ("y", 2)],
            id="dict_with_exprs",
        ),
        pytest.param(
            ("x", 1, "y", 2),
            [("x", 1), ("y", 2)],
            id="variadic_pairs",
        ),
        pytest.param(
            (literal("x"), literal(1), literal("y"), literal(2)),
            [("x", 1), ("y", 2)],
            id="variadic_with_exprs",
        ),
    ],
)
def test_make_map(args, expected):
    ctx = SessionContext()
    batch = pa.RecordBatch.from_arrays([pa.array([1])], names=["a"])
    df = ctx.create_dataframe([[batch]])

    result = df.select(f.make_map(*args).alias("m")).collect()[0].column(0)
    assert result[0].as_py() == expected


def test_make_map_from_two_lists():
    ctx = SessionContext()
    batch = pa.RecordBatch.from_arrays(
        [
            pa.array(["k1", "k2", "k3"]),
            pa.array([10, 20, 30]),
        ],
        names=["keys", "vals"],
    )
    df = ctx.create_dataframe([[batch]])

    m = f.make_map([column("keys")], [column("vals")])
    result = df.select(f.map_keys(m).alias("k")).collect()[0].column(0)
    assert result.to_pylist() == [["k1"], ["k2"], ["k3"]]

    result = df.select(f.map_values(m).alias("v")).collect()[0].column(0)
    assert result.to_pylist() == [[10], [20], [30]]


def test_make_map_odd_args_raises():
    with pytest.raises(ValueError, match="make_map expects"):
        f.make_map("x", 1, "y")


def test_make_map_mismatched_lengths():
    with pytest.raises(ValueError, match="same length"):
        f.make_map(["a", "b"], [1])


@pytest.mark.parametrize(
    ("func", "expected"),
    [
        pytest.param(f.map_keys, ["x", "y"], id="map_keys"),
        pytest.param(f.map_values, [1, 2], id="map_values"),
        pytest.param(
            lambda m: f.map_extract(m, literal("x")),
            [1],
            id="map_extract",
        ),
        pytest.param(
            lambda m: f.map_extract(m, literal("z")),
            [None],
            id="map_extract_missing_key",
        ),
        pytest.param(
            f.map_entries,
            [{"key": "x", "value": 1}, {"key": "y", "value": 2}],
            id="map_entries",
        ),
        pytest.param(
            lambda m: f.element_at(m, literal("y")),
            [2],
            id="element_at",
        ),
    ],
)
def test_map_functions(func, expected):
    ctx = SessionContext()
    batch = pa.RecordBatch.from_arrays([pa.array([1])], names=["a"])
    df = ctx.create_dataframe([[batch]])

    m = f.make_map({"x": 1, "y": 2})
    result = df.select(func(m).alias("out")).collect()[0].column(0)
    assert result[0].as_py() == expected


@pytest.mark.parametrize(
    ("function", "expected_result"),
    [
        (
            f.ascii(column("a")),
            pa.array([72, 87, 33], type=pa.int32()),
        ),  # H = 72; W = 87; ! = 33
        (
            f.bit_length(column("a").cast(pa.string())),
            pa.array([40, 40, 8], type=pa.int32()),
        ),
        (
            f.btrim(literal(" World ")),
            pa.array(["World", "World", "World"], type=pa.string_view()),
        ),
        (f.character_length(column("a")), pa.array([5, 5, 1], type=pa.int32())),
        (f.chr(literal(68)), pa.array(["D", "D", "D"])),
        (
            f.concat_ws("-", column("a"), literal("test")),
            pa.array(["Hello-test", "World-test", "!-test"], type=pa.string_view()),
        ),
        (
            f.concat(column("a").cast(pa.string()), literal("?")),
            pa.array(["Hello?", "World?", "!?"], type=pa.string_view()),
        ),
        (
            f.initcap(column("c")),
            pa.array(["Hello ", " World ", " !"], type=pa.string_view()),
        ),
        (
            f.left(column("a"), literal(3)),
            pa.array(["Hel", "Wor", "!"], type=pa.string_view()),
        ),
        (f.length(column("c")), pa.array([6, 7, 2], type=pa.int32())),
        (
            f.lower(column("a")),
            pa.array(["hello", "world", "!"], type=pa.string_view()),
        ),
        (f.lpad(column("a"), literal(7)), pa.array(["  Hello", "  World", "      !"])),
        (
            f.ltrim(column("c")),
            pa.array(["hello ", "world ", "!"], type=pa.string_view()),
        ),
        (
            f.md5(column("a")),
            pa.array(
                [
                    "8b1a9953c4611296a827abf8c47804d7",
                    "f5a7924e621e84c9280a9a27e1bcb7f6",
                    "9033e0e305f247c0c3c80d0c7848c8b3",
                ],
                type=pa.string_view(),
            ),
        ),
        (f.octet_length(column("a")), pa.array([5, 5, 1], type=pa.int32())),
        (
            f.repeat(column("a"), literal(2)),
            pa.array(["HelloHello", "WorldWorld", "!!"], type=pa.string_view()),
        ),
        (
            f.replace(column("a"), literal("l"), literal("?")),
            pa.array(["He??o", "Wor?d", "!"]),
        ),
        (
            f.reverse(column("a")),
            pa.array(["olleH", "dlroW", "!"], type=pa.string_view()),
        ),
        (
            f.right(column("a"), literal(4)),
            pa.array(["ello", "orld", "!"], type=pa.string_view()),
        ),
        (
            f.rpad(column("a"), literal(8)),
            pa.array(["Hello   ", "World   ", "!       "]),
        ),
        (
            f.rtrim(column("c")),
            pa.array(["hello", " world", " !"], type=pa.string_view()),
        ),
        (
            f.split_part(column("a"), literal("l"), literal(1)),
            pa.array(["He", "Wor", "!"], type=pa.string_view()),
        ),
        (f.contains(column("a"), literal("ell")), pa.array([True, False, False])),
        (f.starts_with(column("a"), literal("Wor")), pa.array([False, True, False])),
        (f.strpos(column("a"), literal("o")), pa.array([5, 2, 0], type=pa.int32())),
        (
            f.substr(column("a"), literal(3)),
            pa.array(["llo", "rld", ""], type=pa.string_view()),
        ),
        (
            f.translate(column("a"), literal("or"), literal("ld")),
            pa.array(["Helll", "Wldld", "!"], type=pa.string_view()),
        ),
        (f.trim(column("c")), pa.array(["hello", "world", "!"], type=pa.string_view())),
        (
            f.upper(column("c")),
            pa.array(["HELLO ", " WORLD ", " !"], type=pa.string_view()),
        ),
        (f.ends_with(column("a"), literal("llo")), pa.array([True, False, False])),
        (
            f.overlay(column("a"), literal("--"), literal(2)),
            pa.array(["H--lo", "W--ld", "!--"]),
        ),
        (
            f.regexp_like(column("a"), literal("(ell|orl)")),
            pa.array([True, True, False]),
        ),
        (
            f.regexp_match(column("a"), literal("(ell|orl)")),
            pa.array([["ell"], ["orl"], None], type=pa.list_(pa.string_view())),
        ),
        (
            f.regexp_replace(column("a"), literal("(ell|orl)"), literal("-")),
            pa.array(["H-o", "W-d", "!"], type=pa.string_view()),
        ),
        (
            f.regexp_count(column("a"), literal("(ell|orl)"), start=literal(1)),
            pa.array([1, 1, 0], type=pa.int64()),
        ),
        (
            f.regexp_count(column("a"), literal("(ell|orl)")),
            pa.array([1, 1, 0], type=pa.int64()),
        ),
        (
            f.regexp_instr(column("a"), literal("(ell|orl)")),
            pa.array([2, 2, 0], type=pa.int64()),
        ),
        (
            f.regexp_instr(column("a"), literal("([lr])"), n=literal(2)),
            pa.array([4, 4, 0], type=pa.int64()),
        ),
        (
            f.regexp_instr(
                column("a"),
                literal("(x)?([hw])"),
                start=literal(1),
                n=literal(1),
                flags=literal("i"),
                sub_expr=literal(2),
            ),
            pa.array([1, 1, 0], type=pa.int64()),
        ),
        (
            f.regexp_instr(column("a"), literal("([hw])"), flags=literal("i")),
            pa.array([1, 1, 0], type=pa.int64()),
        ),
        (
            f.regexp_instr(column("a"), literal("(x)?([HW])"), sub_expr=literal(2)),
            pa.array([1, 1, 0], type=pa.int64()),
        ),
    ],
)
def test_string_functions(df, function, expected_result):
    df = df.select(function)
    result = df.collect()
    assert len(result) == 1
    result = result[0]
    assert result.column(0) == expected_result


def test_hash_functions(df):
    exprs = [
        f.digest(column("a"), m)
        for m in (
            "md5",
            "sha224",
            "sha256",
            "sha384",
            "sha512",
            "blake2s",
            "blake3",
        )
    ]
    df = df.select(
        *exprs,
        f.sha224(column("a")),
        f.sha256(column("a")),
        f.sha384(column("a")),
        f.sha512(column("a")),
    )
    result = df.collect()
    assert len(result) == 1
    result = result[0]
    b = bytearray.fromhex
    assert result.column(0) == pa.array(
        [
            b("8B1A9953C4611296A827ABF8C47804D7"),
            b("F5A7924E621E84C9280A9A27E1BCB7F6"),
            b("9033E0E305F247C0C3C80D0C7848C8B3"),
        ]
    )
    assert result.column(1) == pa.array(
        [
            b("4149DA18AA8BFC2B1E382C6C26556D01A92C261B6436DAD5E3BE3FCC"),
            b("12972632B6D3B6AA52BD6434552F08C1303D56B817119406466E9236"),
            b("6641A7E8278BCD49E476E7ACAE158F4105B2952D22AEB2E0B9A231A0"),
        ]
    )
    assert result.column(2) == pa.array(
        [
            b("185F8DB32271FE25F561A6FC938B2E264306EC304EDA518007D1764826381969"),
            b("78AE647DC5544D227130A0682A51E30BC7777FBB6D8A8F17007463A3ECD1D524"),
            b("BB7208BC9B5D7C04F1236A82A0093A5E33F40423D5BA8D4266F7092C3BA43B62"),
        ]
    )
    assert result.column(3) == pa.array(
        [
            b(
                "3519FE5AD2C596EFE3E276A6F351B8FC"
                "0B03DB861782490D45F7598EBD0AB5FD"
                "5520ED102F38C4A5EC834E98668035FC"
            ),
            b(
                "ED7CED84875773603AF90402E42C65F3"
                "B48A5E77F84ADC7A19E8F3E8D3101010"
                "22F552AEC70E9E1087B225930C1D260A"
            ),
            b(
                "1D0EC8C84EE9521E21F06774DE232367"
                "B64DE628474CB5B2E372B699A1F55AE3"
                "35CC37193EF823E33324DFD9A70738A6"
            ),
        ]
    )
    assert result.column(4) == pa.array(
        [
            b(
                "3615F80C9D293ED7402687F94B22D58E"
                "529B8CC7916F8FAC7FDDF7FBD5AF4CF7"
                "77D3D795A7A00A16BF7E7F3FB9561EE9"
                "BAAE480DA9FE7A18769E71886B03F315"
            ),
            b(
                "8EA77393A42AB8FA92500FB077A9509C"
                "C32BC95E72712EFA116EDAF2EDFAE34F"
                "BB682EFDD6C5DD13C117E08BD4AAEF71"
                "291D8AACE2F890273081D0677C16DF0F"
            ),
            b(
                "3831A6A6155E509DEE59A7F451EB3532"
                "4D8F8F2DF6E3708894740F98FDEE2388"
                "9F4DE5ADB0C5010DFB555CDA77C8AB5D"
                "C902094C52DE3278F35A75EBC25F093A"
            ),
        ]
    )
    assert result.column(5) == pa.array(
        [
            b("F73A5FBF881F89B814871F46E26AD3FA37CB2921C5E8561618639015B3CCBB71"),
            b("B792A0383FB9E7A189EC150686579532854E44B71AC394831DAED169BA85CCC5"),
            b("27988A0E51812297C77A433F635233346AEE29A829DCF4F46E0F58F402C6CFCB"),
        ]
    )
    assert result.column(6) == pa.array(
        [
            b("FBC2B0516EE8744D293B980779178A3508850FDCFE965985782C39601B65794F"),
            b("BF73D18575A736E4037D45F9E316085B86C19BE6363DE6AA789E13DEAACC1C4E"),
            b("C8D11B9F7237E4034ADBCD2005735F9BC4C597C75AD89F4492BEC8F77D15F7EB"),
        ]
    )
    assert result.column(7) == result.column(1)  # SHA-224
    assert result.column(8) == result.column(2)  # SHA-256
    assert result.column(9) == result.column(3)  # SHA-384
    assert result.column(10) == result.column(4)  # SHA-512


def test_temporal_functions(df):
    df = df.select(
        f.date_part("month", column("d")),
        f.datepart("year", column("d")),
        f.date_trunc("month", column("d")),
        f.datetrunc("day", column("d")),
        f.date_bin(
            literal("15 minutes").cast(pa.string()),
            column("d"),
            literal("2001-01-01 00:02:30").cast(pa.string()),
        ),
        f.from_unixtime(literal(1673383974)),
        f.to_timestamp(literal("2023-09-07 05:06:14.523952")),
        f.to_timestamp_seconds(literal("2023-09-07 05:06:14.523952")),
        f.to_timestamp_millis(literal("2023-09-07 05:06:14.523952")),
        f.to_timestamp_micros(literal("2023-09-07 05:06:14.523952")),
        f.extract("day", column("d")),
        f.to_timestamp(
            literal("2023-09-07 05:06:14.523952000"), literal("%Y-%m-%d %H:%M:%S.%f")
        ),
        f.to_timestamp_seconds(
            literal("2023-09-07 05:06:14.523952000"), literal("%Y-%m-%d %H:%M:%S.%f")
        ),
        f.to_timestamp_millis(
            literal("2023-09-07 05:06:14.523952000"), literal("%Y-%m-%d %H:%M:%S.%f")
        ),
        f.to_timestamp_micros(
            literal("2023-09-07 05:06:14.523952000"), literal("%Y-%m-%d %H:%M:%S.%f")
        ),
        f.to_timestamp_nanos(literal("2023-09-07 05:06:14.523952")),
        f.to_timestamp_nanos(
            literal("2023-09-07 05:06:14.523952000"), literal("%Y-%m-%d %H:%M:%S.%f")
        ),
        f.to_time(literal("12:30:45")),
        f.to_time(literal("12-30-45"), literal("%H-%M-%S")),
        f.to_date(literal("2017-05-31")),
        f.to_date(literal("2017-05-31"), literal("%Y-%m-%d")),
        f.to_local_time(column("d")),
        f.to_char(column("d"), literal("%d-%m-%Y")),
    )
    result = df.collect()
    assert len(result) == 1
    result = result[0]
    assert result.column(0) == pa.array([12, 6, 7], type=pa.int32())
    assert result.column(1) == pa.array([2022, 2027, 2020], type=pa.int32())
    assert result.column(2) == pa.array(
        [
            datetime(2022, 12, 1, tzinfo=DEFAULT_TZ),
            datetime(2027, 6, 1, tzinfo=DEFAULT_TZ),
            datetime(2020, 7, 1, tzinfo=DEFAULT_TZ),
        ],
        type=pa.timestamp("us", tz=DEFAULT_TZ),
    )
    assert result.column(3) == pa.array(
        [
            datetime(2022, 12, 31, tzinfo=DEFAULT_TZ),
            datetime(2027, 6, 26, tzinfo=DEFAULT_TZ),
            datetime(2020, 7, 2, tzinfo=DEFAULT_TZ),
        ],
        type=pa.timestamp("us", tz=DEFAULT_TZ),
    )
    assert result.column(4) == pa.array(
        [
            datetime(2022, 12, 30, 23, 47, 30, tzinfo=DEFAULT_TZ),
            datetime(2027, 6, 25, 23, 47, 30, tzinfo=DEFAULT_TZ),
            datetime(2020, 7, 1, 23, 47, 30, tzinfo=DEFAULT_TZ),
        ],
        type=pa.timestamp("ns", tz=DEFAULT_TZ),
    )
    assert result.column(5) == pa.array(
        [datetime(2023, 1, 10, 20, 52, 54, tzinfo=DEFAULT_TZ)] * 3,
        type=pa.timestamp("s"),
    )
    assert result.column(6) == pa.array(
        [datetime(2023, 9, 7, 5, 6, 14, 523952, tzinfo=DEFAULT_TZ)] * 3,
        type=pa.timestamp("ns"),
    )
    assert result.column(7) == pa.array(
        [datetime(2023, 9, 7, 5, 6, 14, tzinfo=DEFAULT_TZ)] * 3, type=pa.timestamp("s")
    )
    assert result.column(8) == pa.array(
        [datetime(2023, 9, 7, 5, 6, 14, 523000, tzinfo=DEFAULT_TZ)] * 3,
        type=pa.timestamp("ms"),
    )
    assert result.column(9) == pa.array(
        [datetime(2023, 9, 7, 5, 6, 14, 523952, tzinfo=DEFAULT_TZ)] * 3,
        type=pa.timestamp("us"),
    )
    assert result.column(10) == pa.array([31, 26, 2], type=pa.int32())
    assert result.column(11) == pa.array(
        [datetime(2023, 9, 7, 5, 6, 14, 523952, tzinfo=DEFAULT_TZ)] * 3,
        type=pa.timestamp("ns"),
    )
    assert result.column(12) == pa.array(
        [datetime(2023, 9, 7, 5, 6, 14, tzinfo=DEFAULT_TZ)] * 3,
        type=pa.timestamp("s"),
    )
    assert result.column(13) == pa.array(
        [datetime(2023, 9, 7, 5, 6, 14, 523000, tzinfo=DEFAULT_TZ)] * 3,
        type=pa.timestamp("ms"),
    )
    assert result.column(14) == pa.array(
        [datetime(2023, 9, 7, 5, 6, 14, 523952, tzinfo=DEFAULT_TZ)] * 3,
        type=pa.timestamp("us"),
    )
    assert result.column(15) == pa.array(
        [datetime(2023, 9, 7, 5, 6, 14, 523952, tzinfo=DEFAULT_TZ)] * 3,
        type=pa.timestamp("ns"),
    )
    assert result.column(16) == pa.array(
        [datetime(2023, 9, 7, 5, 6, 14, 523952, tzinfo=DEFAULT_TZ)] * 3,
        type=pa.timestamp("ns"),
    )
    assert result.column(17) == pa.array(
        [time(12, 30, 45)] * 3,
        type=pa.time64("ns"),
    )
    assert result.column(18) == pa.array(
        [time(12, 30, 45)] * 3,
        type=pa.time64("ns"),
    )
    assert result.column(19) == pa.array(
        [date(2017, 5, 31)] * 3,
        type=pa.date32(),
    )
    assert result.column(20) == pa.array(
        [date(2017, 5, 31)] * 3,
        type=pa.date32(),
    )
    assert result.column(21) == pa.array(
        [
            datetime(2022, 12, 31, tzinfo=DEFAULT_TZ),
            datetime(2027, 6, 26, tzinfo=DEFAULT_TZ),
            datetime(2020, 7, 2, tzinfo=DEFAULT_TZ),
        ],
        type=pa.timestamp("us"),
    )

    assert result.column(22) == pa.array(
        [
            "31-12-2022",
            "26-06-2027",
            "02-07-2020",
        ],
        type=pa.string(),
    )


def test_to_time_invalid_input(df):
    with pytest.raises(Exception, match=r"Error parsing 'not-a-time' as time"):
        df.select(f.to_time(literal("not-a-time"))).collect()


def test_to_time_mismatched_formatter(df):
    with pytest.raises(Exception, match=r"Error parsing '12:30:45' as time"):
        df.select(f.to_time(literal("12:30:45"), literal("%Y-%m-%d"))).collect()


def test_to_date_invalid_input(df):
    with pytest.raises(Exception, match=r"Date32"):
        df.select(f.to_date(literal("not-a-date"))).collect()


def test_temporal_formatter_accepts_bare_str(df):
    result = df.select(
        f.to_time(literal("12h30m45s"), "%Hh%Mm%Ss").alias("t")
    ).collect()[0]
    assert str(result.column(0)[0].as_py()) == "12:30:45"


def test_today_returns_date32(df):
    result = df.select(f.today().alias("today")).collect()[0]
    assert result.column(0).type == pa.date32()


def test_today_alias_matches_current_date(df):
    result = df.select(
        f.current_date().alias("current_date"),
        f.today().alias("today"),
    ).collect()[0]

    assert result.column(0) == result.column(1)


def test_current_timestamp_alias_matches_now(df):
    result = df.select(
        f.now().alias("now"),
        f.current_timestamp().alias("current_timestamp"),
    ).collect()[0]

    assert result.column(0) == result.column(1)


def test_date_format_alias_matches_to_char(df):
    result = df.select(
        f.to_char(
            f.to_timestamp(literal("2021-01-01T00:00:00")), literal("%Y/%m/%d")
        ).alias("to_char"),
        f.date_format(
            f.to_timestamp(literal("2021-01-01T00:00:00")), literal("%Y/%m/%d")
        ).alias("date_format"),
    ).collect()[0]

    assert result.column(0) == result.column(1)
    assert result.column(0)[0].as_py() == "2021/01/01"


def test_make_time(df):
    ctx = SessionContext()
    df_time = ctx.from_pydict({"h": [12], "m": [30], "s": [0]})
    result = df_time.select(
        f.make_time(column("h"), column("m"), column("s")).alias("t")
    ).collect()[0]

    assert result.column(0)[0].as_py() == time(12, 30)


@pytest.mark.parametrize("cast_fn", [f.arrow_cast, f.arrow_try_cast])
@pytest.mark.parametrize(
    ("data_type", "expected"),
    [
        ("Float64", pa.array([4.0, 5.0, 6.0], type=pa.float64())),
        ("Int32", pa.array([4, 5, 6], type=pa.int32())),
        (pa.float64(), pa.array([4.0, 5.0, 6.0], type=pa.float64())),
        (pa.int32(), pa.array([4, 5, 6], type=pa.int32())),
        (pa.string(), pa.array(["4", "5", "6"], type=pa.string())),
    ],
)
def test_arrow_cast_variants(df, cast_fn, data_type, expected):
    """arrow_cast / arrow_try_cast accept str and pyarrow target types."""
    result = df.select(cast_fn(column("b"), data_type).alias("c")).collect()[0]
    assert result.column(0) == expected


def test_arrow_try_cast_null_on_failure():
    ctx = SessionContext()
    batch = pa.RecordBatch.from_arrays([pa.array(["1.5", "oops", "3"])], names=["s"])
    df = ctx.create_dataframe([[batch]])

    result = df.select(f.arrow_try_cast(column("s"), "Float64").alias("c")).collect()[0]

    assert result.column(0).to_pylist() == [1.5, None, 3.0]


def test_arrow_field():
    ctx = SessionContext()
    field = pa.field("val", pa.int64(), metadata={"k": "v"})
    schema = pa.schema([field])
    batch = pa.RecordBatch.from_arrays([pa.array([1])], schema=schema)
    df = ctx.create_dataframe([[batch]])

    out = (
        df.select(f.arrow_field(column("val")).alias("f"))
        .collect_column("f")[0]
        .as_py()
    )
    assert out == {
        "name": "val",
        "data_type": "Int64",
        "nullable": True,
        "metadata": [("k", "v")],
    }


@pytest.mark.parametrize(
    ("cast_fn", "values", "expected"),
    [
        (f.cast_to_type, pa.array([4, 5, 6]), [4.0, 5.0, 6.0]),
        (f.try_cast_to_type, pa.array(["oops", "2", "3"]), [None, 2.0, 3.0]),
    ],
)
def test_cast_to_type(cast_fn, values, expected):
    """cast_to_type / try_cast_to_type take target type from ``type_ref``."""
    ctx = SessionContext()
    batch = pa.RecordBatch.from_arrays(
        [values, pa.array([1.0, 2.0, 3.0])], names=["v", "fl"]
    )
    df = ctx.create_dataframe([[batch]])

    result = df.select(cast_fn(column("v"), column("fl")).alias("c")).collect()[0]

    assert result.column(0).to_pylist() == expected
    assert result.column(0).type == pa.float64()


def test_with_metadata_round_trip(df):
    df = df.select(f.with_metadata(column("b"), {"unit": "ms"}).alias("b"))
    result = df.select(f.arrow_metadata(column("b"), "unit").alias("u")).collect_column(
        "u"
    )
    assert result[0].as_py() == "ms"


def test_with_metadata_empty_dict_noop(df):
    out = df.select(f.with_metadata(column("b"), {}).alias("b")).collect()[0]
    assert out.column(0) == pa.array([4, 5, 6])


def test_with_metadata_empty_key_raises():
    with pytest.raises(ValueError, match="non-empty"):
        f.with_metadata(column("b"), {"": "v"})


def test_case(df):
    df = df.select(
        f.case(column("b")).when(literal(4), literal(10)).otherwise(literal(8)),
        f.case(column("a"))
        .when(literal("Hello"), literal("Hola"))
        .when(literal("World"), literal("Mundo"))
        .otherwise(literal("!!")),
        f.case(column("a"))
        .when(literal("Hello"), literal("Hola"))
        .when(literal("World"), literal("Mundo"))
        .end(),
    )

    result = df.collect()
    result = result[0]
    assert result.column(0) == pa.array([10, 8, 8])
    assert result.column(1) == pa.array(["Hola", "Mundo", "!!"], type=pa.string_view())
    assert result.column(2) == pa.array(["Hola", "Mundo", None], type=pa.string_view())


def test_when_with_no_base(df):
    df.show()
    df = df.select(
        column("b"),
        f.when(column("b") > literal(5), literal("too big"))
        .when(column("b") < literal(5), literal("too small"))
        .otherwise(literal("just right"))
        .alias("goldilocks"),
        f.when(column("a") == literal("Hello"), column("a")).end().alias("greeting"),
    )
    df.show()

    result = df.collect()
    result = result[0]
    assert result.column(0) == pa.array([4, 5, 6])
    assert result.column(1) == pa.array(
        ["too small", "just right", "too big"], type=pa.string_view()
    )
    assert result.column(2) == pa.array(["Hello", None, None], type=pa.string_view())


def test_regr_funcs_sql(df):
    # test case base on
    # https://github.com/apache/arrow-datafusion/blob/d1361d56b9a9e0c165d3d71a8df6795d2a5f51dd/datafusion/core/tests/sqllogictests/test_files/aggregate.slt#L2330
    ctx = SessionContext()
    result = ctx.sql(
        "select regr_slope(1,1), regr_intercept(1,1), "
        "regr_count(1,1), regr_r2(1,1), regr_avgx(1,1), "
        "regr_avgy(1,1), regr_sxx(1,1), regr_syy(1,1), "
        "regr_sxy(1,1);"
    ).collect()

    assert result[0].column(0) == pa.array([None], type=pa.float64())
    assert result[0].column(1) == pa.array([None], type=pa.float64())
    assert result[0].column(2) == pa.array([1], type=pa.uint64())
    assert result[0].column(3) == pa.array([None], type=pa.float64())
    assert result[0].column(4) == pa.array([1], type=pa.float64())
    assert result[0].column(5) == pa.array([1], type=pa.float64())
    assert result[0].column(6) == pa.array([0], type=pa.float64())
    assert result[0].column(7) == pa.array([0], type=pa.float64())
    assert result[0].column(8) == pa.array([0], type=pa.float64())


def test_regr_funcs_sql_2():
    # test case based on `regr_*() basic tests
    # https://github.com/apache/datafusion/blob/d1361d56b9a9e0c165d3d71a8df6795d2a5f51dd/datafusion/core/tests/sqllogictests/test_files/aggregate.slt#L2358C1-L2374C1
    ctx = SessionContext()

    # Perform the regression functions using SQL
    result_sql = ctx.sql(
        "select "
        "regr_slope(column2, column1), "
        "regr_intercept(column2, column1), "
        "regr_count(column2, column1), "
        "regr_r2(column2, column1), "
        "regr_avgx(column2, column1), "
        "regr_avgy(column2, column1), "
        "regr_sxx(column2, column1), "
        "regr_syy(column2, column1), "
        "regr_sxy(column2, column1) "
        "from (values (1,2), (2,4), (3,6))"
    ).collect()

    # Assertions for SQL results
    assert result_sql[0].column(0) == pa.array([2], type=pa.float64())
    assert result_sql[0].column(1) == pa.array([0], type=pa.float64())
    assert result_sql[0].column(2) == pa.array([3], type=pa.uint64())
    assert result_sql[0].column(3) == pa.array([1], type=pa.float64())
    assert result_sql[0].column(4) == pa.array([2], type=pa.float64())
    assert result_sql[0].column(5) == pa.array([4], type=pa.float64())
    assert result_sql[0].column(6) == pa.array([2], type=pa.float64())
    assert result_sql[0].column(7) == pa.array([8], type=pa.float64())
    assert result_sql[0].column(8) == pa.array([4], type=pa.float64())


@pytest.mark.parametrize(
    ("func", "expected"),
    [
        pytest.param(f.regr_slope(column("c2"), column("c1")), [4.6], id="regr_slope"),
        pytest.param(
            f.regr_slope(column("c2"), column("c1"), filter=column("c1") > literal(2)),
            [8],
            id="regr_slope_filter",
        ),
        pytest.param(
            f.regr_intercept(column("c2"), column("c1")), [-4], id="regr_intercept"
        ),
        pytest.param(
            f.regr_intercept(
                column("c2"), column("c1"), filter=column("c1") > literal(2)
            ),
            [-16],
            id="regr_intercept_filter",
        ),
        pytest.param(f.regr_count(column("c2"), column("c1")), [4], id="regr_count"),
        pytest.param(
            f.regr_count(column("c2"), column("c1"), filter=column("c1") > literal(2)),
            [2],
            id="regr_count_filter",
        ),
        pytest.param(f.regr_r2(column("c2"), column("c1")), [0.92], id="regr_r2"),
        pytest.param(
            f.regr_r2(column("c2"), column("c1"), filter=column("c1") > literal(2)),
            [1.0],
            id="regr_r2_filter",
        ),
        pytest.param(f.regr_avgx(column("c2"), column("c1")), [2.5], id="regr_avgx"),
        pytest.param(
            f.regr_avgx(column("c2"), column("c1"), filter=column("c1") > literal(2)),
            [3.5],
            id="regr_avgx_filter",
        ),
        pytest.param(f.regr_avgy(column("c2"), column("c1")), [7.5], id="regr_avgy"),
        pytest.param(
            f.regr_avgy(column("c2"), column("c1"), filter=column("c1") > literal(2)),
            [12.0],
            id="regr_avgy_filter",
        ),
        pytest.param(f.regr_sxx(column("c2"), column("c1")), [5.0], id="regr_sxx"),
        pytest.param(
            f.regr_sxx(column("c2"), column("c1"), filter=column("c1") > literal(2)),
            [0.5],
            id="regr_sxx_filter",
        ),
        pytest.param(f.regr_syy(column("c2"), column("c1")), [115.0], id="regr_syy"),
        pytest.param(
            f.regr_syy(column("c2"), column("c1"), filter=column("c1") > literal(2)),
            [32.0],
            id="regr_syy_filter",
        ),
        pytest.param(f.regr_sxy(column("c2"), column("c1")), [23.0], id="regr_sxy"),
        pytest.param(
            f.regr_sxy(column("c2"), column("c1"), filter=column("c1") > literal(2)),
            [4.0],
            id="regr_sxy_filter",
        ),
    ],
)
def test_regr_funcs_df(func, expected):
    # test case based on `regr_*() basic tests
    # https://github.com/apache/datafusion/blob/d1361d56b9a9e0c165d3d71a8df6795d2a5f51dd/datafusion/core/tests/sqllogictests/test_files/aggregate.slt#L2358C1-L2374C1

    ctx = SessionContext()

    # Create a DataFrame
    data = {"c1": [1, 2, 3, 4, 5, None], "c2": [2, 4, 8, 16, None, 64]}
    df = ctx.from_pydict(data, name="test_table")

    # Perform the regression function using DataFrame API
    df = df.aggregate([], [func.alias("result")])
    df.show()

    expected_dict = {
        "result": expected,
    }

    assert df.collect()[0].to_pydict() == expected_dict


def test_binary_string_functions(df):
    df = df.select(
        f.encode(column("a").cast(pa.string()), "base64"),
        f.decode(
            f.encode(column("a").cast(pa.string()), "base64"),
            "base64",
        ),
    )
    result = df.collect()
    assert len(result) == 1
    result = result[0]
    assert result.column(0) == pa.array(["SGVsbG8", "V29ybGQ", "IQ"])
    assert pa.array(result.column(1)).cast(pa.string()) == pa.array(
        ["Hello", "World", "!"]
    )


@pytest.mark.parametrize(
    ("python_datatype", "name", "expected"),
    [
        pytest.param(bool, "e", pa.bool_(), id="bool"),
        pytest.param(int, "b", pa.int64(), id="int"),
        pytest.param(float, "b", pa.float64(), id="float"),
        pytest.param(str, "b", pa.string(), id="str"),
    ],
)
def test_cast(df, python_datatype, name: str, expected):
    df = df.select(
        column(name).cast(python_datatype).alias("actual"),
        column(name).cast(expected).alias("expected"),
    )
    result = df.collect()
    result = result[0]
    assert result.column(0) == result.column(1)


@pytest.mark.parametrize(
    ("negated", "low", "high", "expected"),
    [
        pytest.param(False, 3, 5, {"filtered": [4, 5]}),
        pytest.param(False, 4, 5, {"filtered": [4, 5]}),
        pytest.param(True, 3, 5, {"filtered": [6]}),
        pytest.param(True, 4, 6, []),
    ],
)
def test_between(df, negated, low, high, expected):
    df = df.filter(column("b").between(low, high, negated=negated)).select(
        column("b").alias("filtered")
    )

    actual = df.collect()

    if expected:
        actual = actual[0].to_pydict()
        assert actual == expected
    else:
        assert len(actual) == 0  # the rows are empty


def test_between_default(df):
    df = df.filter(column("b").between(3, 5)).select(column("b").alias("filtered"))
    expected = {"filtered": [4, 5]}

    actual = df.collect()[0].to_pydict()
    assert actual == expected


def test_alias_with_metadata(df):
    df = df.select(f.alias(f.col("a"), "b", {"key": "value"}))
    assert df.schema().field("b").metadata == {b"key": b"value"}


@pytest.fixture
def df_with_nulls():
    ctx = SessionContext()
    # Rows:
    #   0: both values present
    #   1: a/d/h/k null, b/e/i/l present
    #   2: a/d/h/k present, b/e/i/l null
    #   3: all null
    batch = pa.RecordBatch.from_arrays(
        [
            pa.array([1, None, 3, None], type=pa.int64()),
            pa.array([5, 10, None, None], type=pa.int64()),
            pa.array([20, 30, 40, None], type=pa.int64()),
            pa.array(["apple", None, "cherry", None], type=pa.utf8()),
            pa.array(["banana", "date", None, None], type=pa.utf8()),
            pa.array(["x", "y", "z", None], type=pa.utf8()),
            pa.array(
                [
                    datetime(2020, 1, 1, tzinfo=DEFAULT_TZ),
                    None,
                    datetime(2025, 6, 15, tzinfo=DEFAULT_TZ),
                    None,
                ],
                type=pa.timestamp("us", tz="UTC"),
            ),
            pa.array(
                [
                    datetime(2022, 7, 4, tzinfo=DEFAULT_TZ),
                    datetime(2023, 12, 25, tzinfo=DEFAULT_TZ),
                    None,
                    None,
                ],
                type=pa.timestamp("us", tz="UTC"),
            ),
            pa.array([True, None, False, None], type=pa.bool_()),
            pa.array([False, True, None, None], type=pa.bool_()),
        ],
        names=["a", "b", "c", "d", "e", "g", "h", "i", "k", "l"],
    )
    return ctx.create_dataframe([[batch]])


@pytest.mark.parametrize(
    ("expr", "expected"),
    [
        pytest.param(
            f.greatest(column("a"), column("b")),
            pa.array([5, 10, 3, None], type=pa.int64()),
            id="greatest_int",
        ),
        pytest.param(
            f.greatest(column("d"), column("e")),
            pa.array(["banana", "date", "cherry", None], type=pa.utf8()),
            id="greatest_str",
        ),
        pytest.param(
            f.least(column("a"), column("b")),
            pa.array([1, 10, 3, None], type=pa.int64()),
            id="least_int",
        ),
        pytest.param(
            f.least(column("d"), column("e")),
            pa.array(["apple", "date", "cherry", None], type=pa.utf8()),
            id="least_str",
        ),
        pytest.param(
            f.coalesce(column("a"), column("b"), column("c")),
            pa.array([1, 10, 3, None], type=pa.int64()),
            id="coalesce_int",
        ),
        pytest.param(
            f.coalesce(column("d"), column("e"), column("g")),
            pa.array(["apple", "date", "cherry", None], type=pa.utf8()),
            id="coalesce_str",
        ),
        pytest.param(
            f.nvl(column("a"), column("c")),
            pa.array([1, 30, 3, None], type=pa.int64()),
            id="nvl_int",
        ),
        pytest.param(
            f.nvl(column("d"), column("g")),
            pa.array(["apple", "y", "cherry", None], type=pa.utf8()),
            id="nvl_str",
        ),
        pytest.param(
            f.ifnull(column("a"), column("c")),
            pa.array([1, 30, 3, None], type=pa.int64()),
            id="ifnull_int",
        ),
        pytest.param(
            f.ifnull(column("d"), column("g")),
            pa.array(["apple", "y", "cherry", None], type=pa.utf8()),
            id="ifnull_str",
        ),
        pytest.param(
            f.nvl2(column("a"), column("b"), column("c")),
            pa.array([5, 30, None, None], type=pa.int64()),
            id="nvl2_int",
        ),
        pytest.param(
            f.nvl2(column("d"), column("e"), column("g")),
            pa.array(["banana", "y", None, None], type=pa.utf8()),
            id="nvl2_str",
        ),
        pytest.param(
            f.nullif(column("a"), column("b")),
            pa.array([1, None, 3, None], type=pa.int64()),
            id="nullif_int",
        ),
        pytest.param(
            f.nullif(column("d"), column("e")),
            pa.array(["apple", None, "cherry", None], type=pa.utf8()),
            id="nullif_str",
        ),
        pytest.param(
            f.nullif(column("a"), literal(1)),
            pa.array([None, None, 3, None], type=pa.int64()),
            id="nullif_equal_values",
        ),
        pytest.param(
            f.greatest(column("a"), column("b"), column("c")),
            pa.array([20, 30, 40, None], type=pa.int64()),
            id="greatest_variadic",
        ),
        pytest.param(
            f.least(column("a"), column("b"), column("c")),
            pa.array([1, 10, 3, None], type=pa.int64()),
            id="least_variadic",
        ),
        pytest.param(
            f.greatest(column("a"), literal(2)),
            pa.array([2, 2, 3, 2], type=pa.int64()),
            id="greatest_literal",
        ),
        pytest.param(
            f.least(column("a"), literal(2)),
            pa.array([1, 2, 2, 2], type=pa.int64()),
            id="least_literal",
        ),
        pytest.param(
            f.coalesce(column("a"), literal(0)),
            pa.array([1, 0, 3, 0], type=pa.int64()),
            id="coalesce_literal_int",
        ),
        pytest.param(
            f.coalesce(column("d"), literal("default")),
            pa.array(["apple", "default", "cherry", "default"], type=pa.string_view()),
            id="coalesce_literal_str",
        ),
        pytest.param(
            f.nvl(column("a"), literal(99)),
            pa.array([1, 99, 3, 99], type=pa.int64()),
            id="nvl_literal",
        ),
        pytest.param(
            f.ifnull(column("d"), literal("unknown")),
            pa.array(["apple", "unknown", "cherry", "unknown"], type=pa.string_view()),
            id="ifnull_literal",
        ),
        pytest.param(
            f.nvl2(column("a"), literal(1), literal(0)),
            pa.array([1, 0, 1, 0], type=pa.int64()),
            id="nvl2_literal",
        ),
        pytest.param(
            f.greatest(column("h"), column("i")),
            pa.array(
                [
                    datetime(2022, 7, 4, tzinfo=DEFAULT_TZ),
                    datetime(2023, 12, 25, tzinfo=DEFAULT_TZ),
                    datetime(2025, 6, 15, tzinfo=DEFAULT_TZ),
                    None,
                ],
                type=pa.timestamp("us", tz="UTC"),
            ),
            id="greatest_datetime",
        ),
        pytest.param(
            f.least(column("h"), column("i")),
            pa.array(
                [
                    datetime(2020, 1, 1, tzinfo=DEFAULT_TZ),
                    datetime(2023, 12, 25, tzinfo=DEFAULT_TZ),
                    datetime(2025, 6, 15, tzinfo=DEFAULT_TZ),
                    None,
                ],
                type=pa.timestamp("us", tz="UTC"),
            ),
            id="least_datetime",
        ),
        pytest.param(
            f.coalesce(column("h"), column("i")),
            pa.array(
                [
                    datetime(2020, 1, 1, tzinfo=DEFAULT_TZ),
                    datetime(2023, 12, 25, tzinfo=DEFAULT_TZ),
                    datetime(2025, 6, 15, tzinfo=DEFAULT_TZ),
                    None,
                ],
                type=pa.timestamp("us", tz="UTC"),
            ),
            id="coalesce_datetime",
        ),
        pytest.param(
            f.nvl(column("k"), column("l")),
            pa.array([True, True, False, None], type=pa.bool_()),
            id="nvl_bool",
        ),
        pytest.param(
            f.coalesce(column("k"), column("l")),
            pa.array([True, True, False, None], type=pa.bool_()),
            id="coalesce_bool",
        ),
        pytest.param(
            f.nvl2(column("k"), column("k"), column("l")),
            pa.array([True, True, False, None], type=pa.bool_()),
            id="nvl2_bool",
        ),
        pytest.param(
            f.coalesce(
                column("h"),
                literal(datetime(2000, 1, 1, tzinfo=DEFAULT_TZ)),
            ),
            pa.array(
                [
                    datetime(2020, 1, 1, tzinfo=DEFAULT_TZ),
                    datetime(2000, 1, 1, tzinfo=DEFAULT_TZ),
                    datetime(2025, 6, 15, tzinfo=DEFAULT_TZ),
                    datetime(2000, 1, 1, tzinfo=DEFAULT_TZ),
                ],
                type=pa.timestamp("us", tz="UTC"),
            ),
            id="coalesce_literal_datetime",
        ),
        pytest.param(
            f.coalesce(column("k"), literal(value=False)),
            pa.array([True, False, False, False], type=pa.bool_()),
            id="coalesce_literal_bool",
        ),
        pytest.param(
            f.coalesce(column("a"), literal(None), literal(99)),
            pa.array([1, 99, 3, 99], type=pa.int64()),
            id="coalesce_skip_null_literal",
        ),
    ],
)
def test_conditional_functions(df_with_nulls, expr, expected):
    result = df_with_nulls.select(expr.alias("result")).collect()[0]
    assert result.column(0) == expected


@pytest.mark.parametrize(
    ("func", "filter_expr", "expected"),
    [
        (f.percentile_cont, None, 3.0),
        (f.percentile_cont, column("a") > literal(1.0), 3.5),
        (f.quantile_cont, None, 3.0),
    ],
    ids=["no_filter", "with_filter", "quantile_cont_alias"],
)
def test_percentile_cont(func, filter_expr, expected):
    ctx = SessionContext()
    df = ctx.from_pydict({"a": [1.0, 2.0, 3.0, 4.0, 5.0]})
    result = df.aggregate(
        [], [func(column("a"), 0.5, filter=filter_expr).alias("v")]
    ).collect()[0]
    assert result.column(0)[0].as_py() == expected


@pytest.mark.parametrize(
    ("grouping_set_expr", "expected_grouping", "expected_sums"),
    [
        (GroupingSet.rollup(column("a")), [0, 0, 1], [30, 30, 60]),
        (GroupingSet.cube(column("a")), [0, 0, 1], [30, 30, 60]),
        (GroupingSet.rollup("a"), [0, 0, 1], [30, 30, 60]),
        (GroupingSet.cube("a"), [0, 0, 1], [30, 30, 60]),
    ],
    ids=["rollup", "cube", "rollup_str", "cube_str"],
)
def test_grouping_set_single_column(
    grouping_set_expr, expected_grouping, expected_sums
):
    ctx = SessionContext()
    df = ctx.from_pydict({"a": [1, 1, 2], "b": [10, 20, 30]})
    result = df.aggregate(
        [grouping_set_expr],
        [f.sum(column("b")).alias("s"), f.grouping(column("a"))],
    ).sort(column("a").sort(ascending=True, nulls_first=False))
    batches = result.collect()
    g = pa.concat_arrays([b.column(2) for b in batches]).to_pylist()
    s = pa.concat_arrays([b.column("s") for b in batches]).to_pylist()
    assert g == expected_grouping
    assert s == expected_sums


@pytest.mark.parametrize(
    ("grouping_set_expr", "expected_rows"),
    [
        # rollup(a, b) => (a,b), (a), () => 3 + 2 + 1 = 6
        (GroupingSet.rollup(column("a"), column("b")), 6),
        # cube(a, b) => (a,b), (a), (b), () => 3 + 2 + 2 + 1 = 8
        (GroupingSet.cube(column("a"), column("b")), 8),
        (GroupingSet.rollup("a", "b"), 6),
        (GroupingSet.cube("a", "b"), 8),
    ],
    ids=["rollup", "cube", "rollup_str", "cube_str"],
)
def test_grouping_set_multi_column(grouping_set_expr, expected_rows):
    ctx = SessionContext()
    df = ctx.from_pydict({"a": [1, 1, 2], "b": ["x", "y", "x"], "c": [10, 20, 30]})
    result = df.aggregate(
        [grouping_set_expr],
        [f.sum(column("c")).alias("s")],
    )
    total_rows = sum(b.num_rows for b in result.collect())
    assert total_rows == expected_rows


@pytest.mark.parametrize(
    "grouping_set_expr",
    [
        GroupingSet.grouping_sets([column("a")], [column("b")]),
        GroupingSet.grouping_sets(["a"], ["b"]),
    ],
    ids=["expr", "str"],
)
def test_grouping_sets_explicit(grouping_set_expr):
    # Each row's grouping() value tells you which columns are aggregated across.
    ctx = SessionContext()
    df = ctx.from_pydict({"a": ["x", "x", "y"], "b": ["m", "n", "m"], "c": [1, 2, 3]})
    result = df.aggregate(
        [grouping_set_expr],
        [
            f.sum(column("c")).alias("s"),
            f.grouping(column("a")),
            f.grouping(column("b")),
        ],
    ).sort(
        column("a").sort(ascending=True, nulls_first=False),
        column("b").sort(ascending=True, nulls_first=False),
    )
    batches = result.collect()
    ga = pa.concat_arrays([b.column(3) for b in batches]).to_pylist()
    gb = pa.concat_arrays([b.column(4) for b in batches]).to_pylist()
    # Rows grouped by (a): ga=0 (a is a key), gb=1 (b is aggregated across)
    # Rows grouped by (b): ga=1 (a is aggregated across), gb=0 (b is a key)
    assert ga == [0, 0, 1, 1]
    assert gb == [1, 1, 0, 0]


def test_var_population():
    ctx = SessionContext()
    df = ctx.from_pydict({"a": [-1.0, 0.0, 2.0]})
    result = df.aggregate([], [f.var_population(column("a")).alias("v")]).collect()[0]
    # var_population is an alias for var_pop
    expected = df.aggregate([], [f.var_pop(column("a")).alias("v")]).collect()[0]
    assert abs(result.column(0)[0].as_py() - expected.column(0)[0].as_py()) < 1e-10


def test_get_field(df):
    df = df.with_column(
        "s",
        f.named_struct(
            [
                ("x", column("a")),
                ("y", column("b")),
            ]
        ),
    )
    result = df.select(
        f.get_field(column("s"), "x").alias("x_val"),
        f.get_field(column("s"), "y").alias("y_val"),
    ).collect()[0]

    assert result.column(0) == pa.array(["Hello", "World", "!"], type=pa.string_view())
    assert result.column(1) == pa.array([4, 5, 6])


def test_get_field_path(df):
    df = df.with_column(
        "outer",
        f.named_struct(
            [
                (
                    "inner",
                    f.named_struct(
                        [
                            ("x", column("a")),
                            ("y", column("b")),
                        ]
                    ),
                ),
            ]
        ),
    )
    result = df.select(
        f.get_field(column("outer"), "inner", "x").alias("x_val"),
        f.get_field(column("outer"), "inner", "y").alias("y_val"),
    ).collect()[0]

    assert result.column(0) == pa.array(["Hello", "World", "!"], type=pa.string_view())
    assert result.column(1) == pa.array([4, 5, 6])


def test_get_field_requires_a_name():
    with pytest.raises(ValueError, match="at least one field name"):
        f.get_field(column("s"))


def test_arrow_metadata():
    ctx = SessionContext()
    field = pa.field("val", pa.int64(), metadata={"key1": "value1", "key2": "value2"})
    schema = pa.schema([field])
    batch = pa.RecordBatch.from_arrays([pa.array([1, 2, 3])], schema=schema)
    df = ctx.create_dataframe([[batch]])

    # One-argument form: returns a Map of all metadata key-value pairs
    result = df.select(
        f.arrow_metadata(column("val")).alias("meta"),
    ).collect()[0]
    assert result.column(0).type == pa.map_(pa.utf8(), pa.utf8())
    meta = result.column(0)[0].as_py()
    assert ("key1", "value1") in meta
    assert ("key2", "value2") in meta

    # Two-argument form: returns the value for a specific metadata key
    result = df.select(
        f.arrow_metadata(column("val"), "key1").alias("meta_val"),
    ).collect()[0]
    assert result.column(0)[0].as_py() == "value1"


def test_version():
    ctx = SessionContext()
    df = ctx.from_pydict({"a": [1]})
    result = df.select(f.version().alias("v")).collect()[0]
    version_str = result.column(0)[0].as_py()
    assert "Apache DataFusion" in version_str


def test_row(df):
    result = df.select(
        f.row(column("a"), column("b")).alias("r"),
        f.struct(column("a"), column("b")).alias("s"),
    ).collect()[0]
    # row is an alias for struct, so they should produce the same output
    assert result.column(0) == result.column(1)


def test_union_tag():
    ctx = SessionContext()
    types = pa.array([0, 1, 0], type=pa.int8())
    offsets = pa.array([0, 0, 1], type=pa.int32())
    children = [pa.array([1, 2]), pa.array(["hello"])]
    arr = pa.UnionArray.from_dense(types, offsets, children, ["int", "str"], [0, 1])
    df = ctx.create_dataframe([[pa.RecordBatch.from_arrays([arr], names=["u"])]])

    result = df.select(f.union_tag(column("u")).alias("tag")).collect()[0]
    assert result.column(0).to_pylist() == ["int", "str", "int"]


def test_union_extract():
    ctx = SessionContext()
    types = pa.array([0, 1, 0], type=pa.int8())
    offsets = pa.array([0, 0, 1], type=pa.int32())
    children = [pa.array([1, 2]), pa.array(["hello"])]
    arr = pa.UnionArray.from_dense(types, offsets, children, ["int", "str"], [0, 1])
    df = ctx.create_dataframe([[pa.RecordBatch.from_arrays([arr], names=["u"])]])

    result = df.select(f.union_extract(column("u"), "int").alias("val")).collect()[0]
    assert result.column(0).to_pylist() == [1, None, 2]


@pytest.mark.parametrize("func", [f.array_any_value, f.list_any_value])
def test_any_value_aliases(func):
    ctx = SessionContext()
    df = ctx.from_pydict({"a": [[None, 2, 3], [None, None, None], [1, 2, 3]]})
    result = df.select(func(column("a")).alias("v")).collect()
    values = [row.as_py() for row in result[0].column(0)]
    assert values[0] == 2
    assert values[1] is None
    assert values[2] == 1


@pytest.mark.parametrize("func", [f.array_distance, f.list_distance])
def test_array_distance_aliases(func):
    ctx = SessionContext()
    df = ctx.from_pydict({"a": [[1.0, 2.0]], "b": [[1.0, 4.0]]})
    result = df.select(func(column("a"), column("b")).alias("v")).collect()
    assert result[0].column(0)[0].as_py() == pytest.approx(2.0)


@pytest.mark.parametrize(
    ("func", "expected"),
    [
        (f.array_max, [5, 10]),
        (f.list_max, [5, 10]),
        (f.array_min, [1, 2]),
        (f.list_min, [1, 2]),
    ],
)
def test_array_min_max(func, expected):
    ctx = SessionContext()
    df = ctx.from_pydict({"a": [[1, 5, 3], [10, 2]]})
    result = df.select(func(column("a")).alias("v")).collect()
    values = [row.as_py() for row in result[0].column(0)]
    assert values == expected


@pytest.mark.parametrize("func", [f.array_reverse, f.list_reverse])
def test_array_reverse_aliases(func):
    ctx = SessionContext()
    df = ctx.from_pydict({"a": [[1, 2, 3], [4, 5]]})
    result = df.select(func(column("a")).alias("v")).collect()
    values = [row.as_py() for row in result[0].column(0)]
    assert values == [[3, 2, 1], [5, 4]]


@pytest.mark.parametrize("func", [f.arrays_zip, f.list_zip])
def test_arrays_zip_aliases(func):
    ctx = SessionContext()
    df = ctx.from_pydict({"a": [[1, 2]], "b": [[3, 4]]})
    result = df.select(func(column("a"), column("b")).alias("v")).collect()
    values = result[0].column(0)[0].as_py()
    assert values == [{"1": 1, "2": 3}, {"1": 2, "2": 4}]


@pytest.mark.parametrize("func", [f.string_to_array, f.string_to_list])
def test_string_to_array_aliases(func):
    ctx = SessionContext()
    df = ctx.from_pydict({"a": ["hello,world,foo"]})
    result = df.select(func(column("a"), literal(",")).alias("v")).collect()
    assert result[0].column(0)[0].as_py() == ["hello", "world", "foo"]


def test_string_to_array_with_null_string():
    ctx = SessionContext()
    df = ctx.from_pydict({"a": ["hello,NA,world"]})
    result = df.select(
        f.string_to_array(column("a"), literal(","), literal("NA")).alias("v")
    ).collect()
    values = result[0].column(0)[0].as_py()
    assert values == ["hello", None, "world"]


@pytest.mark.parametrize("func", [f.gen_series, f.generate_series])
def test_gen_series_aliases(func):
    ctx = SessionContext()
    df = ctx.from_pydict({"a": [0]})
    result = df.select(func(literal(1), literal(5)).alias("v")).collect()
    assert result[0].column(0)[0].as_py() == [1, 2, 3, 4, 5]


def test_gen_series_with_step():
    ctx = SessionContext()
    df = ctx.from_pydict({"a": [0]})
    result = df.select(
        f.gen_series(literal(1), literal(10), literal(3)).alias("v")
    ).collect()
    assert result[0].column(0)[0].as_py() == [1, 4, 7, 10]


class TestPythonicNativeTypes:
    """Tests for accepting native Python types instead of requiring lit()."""

    def test_split_part_native(self):
        ctx = SessionContext()
        df = ctx.from_pydict({"a": ["a,b,c"]})
        result = df.select(f.split_part(column("a"), ",", 2).alias("s")).collect()
        assert result[0].column(0)[0].as_py() == "b"

    def test_encode_native_str(self):
        ctx = SessionContext()
        df = ctx.from_pydict({"a": ["hello"]})
        result = df.select(f.encode(column("a"), "base64").alias("e")).collect()
        assert result[0].column(0)[0].as_py() == "aGVsbG8"

    def test_date_part_native_str(self):
        ctx = SessionContext()
        df = ctx.from_pydict({"a": ["2021-07-15T00:00:00"]})
        df = df.select(f.to_timestamp(column("a")).alias("a"))
        with warnings.catch_warnings():
            warnings.simplefilter("error", DeprecationWarning)
            result = df.select(f.date_part("year", column("a")).alias("y")).collect()
        assert result[0].column(0)[0].as_py() == 2021

    @pytest.mark.parametrize(
        ("func", "name"),
        [
            pytest.param(f.date_part, "date_part", id="date_part"),
            pytest.param(f.datepart, "datepart", id="datepart"),
            pytest.param(f.extract, "extract", id="extract"),
        ],
    )
    def test_date_part_expr_part_warns_deprecated(self, func, name):
        with pytest.warns(
            DeprecationWarning,
            match=rf"Passing Expr for {name}\(\) argument 'part' is deprecated",
        ):
            expr = func(literal("year"), column("a"))
        assert expr is not None

    def test_date_trunc_native_str(self):
        ctx = SessionContext()
        df = ctx.from_pydict({"a": ["2021-07-15T12:34:56"]})
        df = df.select(f.to_timestamp(column("a")).alias("a"))
        with warnings.catch_warnings():
            warnings.simplefilter("error", DeprecationWarning)
            result = df.select(f.date_trunc("month", column("a")).alias("t")).collect()
        assert str(result[0].column(0)[0].as_py()) == "2021-07-01 00:00:00"

    @pytest.mark.parametrize(
        ("func", "name"),
        [
            pytest.param(f.date_trunc, "date_trunc", id="date_trunc"),
            pytest.param(f.datetrunc, "datetrunc", id="datetrunc"),
        ],
    )
    def test_date_trunc_expr_part_warns_deprecated(self, func, name):
        with pytest.warns(
            DeprecationWarning,
            match=rf"Passing Expr for {name}\(\) argument 'part' is deprecated",
        ):
            expr = func(literal("month"), column("a"))
        assert expr is not None

    def test_left_native_int(self):
        ctx = SessionContext()
        df = ctx.from_pydict({"a": ["the cat"]})
        result = df.select(f.left(column("a"), 3).alias("l")).collect()
        assert result[0].column(0)[0].as_py() == "the"

    def test_round_native_int(self):
        ctx = SessionContext()
        df = ctx.from_pydict({"a": [1.567]})
        result = df.select(f.round(column("a"), 2).alias("r")).collect()
        assert result[0].column(0)[0].as_py() == 1.57

    def test_regexp_count_native(self):
        ctx = SessionContext()
        df = ctx.from_pydict({"a": ["abcabc"]})
        result = df.select(
            f.regexp_count(column("a"), "abc", start=4, flags="i").alias("c")
        ).collect()
        assert result[0].column(0)[0].as_py() == 1

    def test_log_native_int(self):
        ctx = SessionContext()
        df = ctx.from_pydict({"a": [100.0]})
        result = df.select(f.log(10, column("a")).alias("l")).collect()
        assert result[0].column(0)[0].as_py() == 2.0

    def test_power_native_int(self):
        ctx = SessionContext()
        df = ctx.from_pydict({"a": [2.0]})
        result = df.select(f.power(column("a"), 3).alias("p")).collect()
        assert result[0].column(0)[0].as_py() == 8.0

    def test_array_slice_native(self):
        ctx = SessionContext()
        df = ctx.from_pydict({"a": [[1, 2, 3, 4]]})
        result = df.select(f.array_slice(column("a"), 2, 3).alias("s")).collect()
        assert result[0].column(0)[0].as_py() == [2, 3]

    def test_string_to_array_native(self):
        ctx = SessionContext()
        df = ctx.from_pydict({"a": ["hello,NA,world"]})
        result = df.select(
            f.string_to_array(column("a"), ",", null_string="NA").alias("v")
        ).collect()
        assert result[0].column(0)[0].as_py() == ["hello", None, "world"]

    def test_regexp_replace_native(self):
        ctx = SessionContext()
        df = ctx.from_pydict({"a": ["a1 b2 c3"]})
        result = df.select(
            f.regexp_replace(column("a"), r"\d+", "X", flags="g").alias("r")
        ).collect()
        assert result[0].column(0)[0].as_py() == "aX bX cX"

    @pytest.mark.parametrize(
        ("func", "arg_name", "expr"),
        [
            pytest.param(
                f.encode,
                "encoding",
                lambda: f.encode(column("a"), literal("base64")),
                id="encode-encoding",
            ),
            pytest.param(
                f.decode,
                "encoding",
                lambda: f.decode(column("a"), literal("base64")),
                id="decode-encoding",
            ),
            pytest.param(
                f.digest,
                "method",
                lambda: f.digest(column("a"), literal("sha256")),
                id="digest-method",
            ),
            pytest.param(
                f.arrow_cast,
                "data_type",
                lambda: f.arrow_cast(column("a"), literal("Float64")),
                id="arrow-cast-data-type",
            ),
            pytest.param(
                f.arrow_try_cast,
                "data_type",
                lambda: f.arrow_try_cast(column("a"), literal("Float64")),
                id="arrow-try-cast-data-type",
            ),
            pytest.param(
                f.arrow_metadata,
                "key",
                lambda: f.arrow_metadata(column("a"), literal("k")),
                id="arrow-metadata-key",
            ),
        ],
    )
    def test_literal_only_expr_args_warn_deprecated(self, func, arg_name, expr):
        with pytest.warns(
            DeprecationWarning,
            match=(
                rf"Passing Expr for {func.__name__}\(\) argument "
                rf"'{arg_name}' is deprecated"
            ),
        ):
            result = expr()
        assert result is not None

    def test_literal_only_native_args_do_not_warn(self):
        with warnings.catch_warnings():
            warnings.simplefilter("error", DeprecationWarning)
            assert f.encode(column("a"), "base64") is not None
            assert f.decode(column("a"), "base64") is not None
            assert f.digest(column("a"), "sha256") is not None
            assert f.arrow_cast(column("a"), "Float64") is not None
            assert f.arrow_try_cast(column("a"), pa.float64()) is not None
            assert f.arrow_metadata(column("a"), "k") is not None

    def test_backward_compat_with_lit(self):
        """Verify that existing code using lit() still works."""
        ctx = SessionContext()
        df = ctx.from_pydict({"a": ["a,b,c"]})
        result = df.select(
            f.split_part(column("a"), literal(","), literal(2)).alias("s")
        ).collect()
        assert result[0].column(0)[0].as_py() == "b"