# 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()),

),