# 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 pyarrow as pa

import pyarrow.dataset as ds

import pytest

from datafusion import Expr, SessionContext, Table, udtf

from datafusion.context import TableProviderExportable

from datafusion.user_defined import TableFunction

def python_table_function_inner(

num_cols: int, num_rows: int, num_batches: int

) -> Table:

rows = list(range(num_rows))

cols = [pa.array(rows) for _ in range(num_cols)]

names = [f"row_{i}" for i in range(num_cols)]

batch = pa.RecordBatch.from_arrays(cols, names=names)

batches = [batch for _ in range(num_batches)]

dataset = ds.dataset(batches)

return Table(dataset)

def test_python_table_function_with_args_class() -> None:

"""Test Python TableFunction using a class with arguments."""

ctx = SessionContext()

class ParameterizedTableFunction:

"""Table function that takes parameters."""

def __call__(self, num_cols: Expr, num_rows: Expr, num_batches: Expr) -> Table:

return python_table_function_inner(

num_cols.to_variant().value_i64(),

num_rows.to_variant().value_i64(),

num_batches.to_variant().value_i64(),

)

# Register the function

table_func = ParameterizedTableFunction()

table_udtf = udtf(table_func, "param_func")

ctx.register_udtf(table_udtf)

# Call with different parameters

result = ctx.sql("SELECT * FROM param_func(5, 10, 2)").collect()

assert len(result) == 2

assert result[0].num_columns == 5

assert result[0].num_rows == 10

def test_python_table_function_decorator() -> None:

"""Test Python TableFunction using decorator syntax."""

ctx = SessionContext()

@udtf("decorated_func")

def my_decorated_func(num_cols: Expr, num_rows: Expr, num_batches: Expr) -> Table:

return python_table_function_inner(

num_cols.to_variant().value_i64(),

num_rows.to_variant().value_i64(),

num_batches.to_variant().value_i64(),

)

ctx.register_udtf(my_decorated_func)

result = ctx.sql("SELECT * FROM decorated_func(3, 7, 2)").collect()

assert len(result) == 2

assert result[0].num_columns == 3

assert result[0].num_rows == 7

def test_python_table_function_no_args() -> None:

"""Test Python TableFunction with no arguments."""

ctx = SessionContext()

@udtf("static_func")

def static_table_func() -> Table:

return python_table_function_inner(2, 3, 1)

ctx.register_udtf(static_table_func)

result = ctx.sql("SELECT * FROM static_func()").collect()

assert len(result) == 1

assert list(result[0].column(0).to_pylist()) == [0, 1, 2]

assert list(result[0].column(1).to_pylist()) == [0, 1, 2]

def test_python_table_function_single_arg() -> None:

"""Test Python TableFunction with a single argument."""

ctx = SessionContext()

@udtf("single_arg_func")

def single_arg_func(n: Expr) -> TableProviderExportable:

return python_table_function_inner(2, n.to_variant().value_i64(), 1)

ctx.register_udtf(single_arg_func)

result = ctx.sql("SELECT * FROM single_arg_func(15)").collect()

assert len(result) == 1

assert result[0].num_columns == 2

assert result[0].num_rows == 15

def test_python_table_function_with_string_args() -> None:

"""Test Python TableFunction with string arguments."""

ctx = SessionContext()

@udtf("string_arg_func")

def string_arg_func(prefix: Expr) -> TableProviderExportable:

prefix_str = prefix.to_variant().value_string()

# Create a table with the prefix in column names

batch = pa.RecordBatch.from_arrays(

[pa.array([1, 2, 3]), pa.array([4, 5, 6])],

names=[f"{prefix_str}_a", f"{prefix_str}_b"],

)

return Table(ds.dataset([batch]))

ctx.register_udtf(string_arg_func)

result = ctx.sql("SELECT * FROM string_arg_func('test')").collect()

assert len(result) == 1

assert result[0].schema.names == ["test_a", "test_b"]

def test_python_table_function_receives_session() -> None:

"""A UDTF registered ``with_session=True`` gets the calling ctx."""

ctx = SessionContext()

captured: list[SessionContext] = []

@udtf("session_aware_func", with_session=True)

def session_aware_func(*, session: SessionContext) -> TableProviderExportable:

captured.append(session)

batch = pa.RecordBatch.from_pydict({"a": [1, 2, 3]})

return Table(ds.dataset([batch]))

ctx.register_udtf(session_aware_func)

result = ctx.sql("SELECT * FROM session_aware_func()").collect()

assert len(captured) == 1

assert isinstance(captured[0], SessionContext)

# Sharing the same catalog confirms the wrapper points at the caller's state.

assert captured[0].catalog().schema().names() == ctx.catalog().schema().names()

assert result[0].column(0).to_pylist() == [1, 2, 3]

def test_python_table_function_session_used_for_metadata() -> None:

"""The UDTF can inspect session state through the passed-in context."""

ctx = SessionContext()

base_batch = pa.RecordBatch.from_pydict({"x": [10, 20, 30]})

ctx.register_batch("base_tbl", base_batch)

seen_tables: list[set[str]] = []

@udtf("table_inventory", with_session=True)

def table_inventory(*, session: SessionContext) -> TableProviderExportable:

# Stash the visible tables to verify the session wired through.

seen_tables.append(session.catalog().schema().names())

batch = pa.RecordBatch.from_pydict({"name": ["base_tbl"]})

return Table(ds.dataset([batch]))

ctx.register_udtf(table_inventory)

result = ctx.sql("SELECT * FROM table_inventory()").collect()

assert seen_tables == [{"base_tbl"}]

assert result[0].column(0).to_pylist() == ["base_tbl"]

def test_python_table_function_class_callable_with_session() -> None:

"""Class-based UDTFs opt in via ``with_session=True``."""

ctx = SessionContext()

captured: list[SessionContext] = []

class SessionAware:

def __call__(

self, n: Expr, *, session: SessionContext

) -> TableProviderExportable:

captured.append(session)

count = n.to_variant().value_i64()

batch = pa.RecordBatch.from_pydict({"a": list(range(count))})

return Table(ds.dataset([batch]))

ctx.register_udtf(udtf(SessionAware(), "session_class_func", with_session=True))

result = ctx.sql("SELECT * FROM session_class_func(3)").collect()

assert len(captured) == 1

assert isinstance(captured[0], SessionContext)

assert result[0].column(0).to_pylist() == [0, 1, 2]

def test_python_table_function_without_session_flag_no_injection() -> None:

"""Default registration (no ``with_session``) calls func positionally."""

ctx = SessionContext()

@udtf("plain_func")

def plain_func(n: Expr) -> TableProviderExportable:

count = n.to_variant().value_i64()

batch = pa.RecordBatch.from_pydict({"a": list(range(count))})

return Table(ds.dataset([batch]))

ctx.register_udtf(plain_func)

result = ctx.sql("SELECT * FROM plain_func(4)").collect()

assert result[0].column(0).to_pylist() == [0, 1, 2, 3]

def test_with_session_rejected_for_ffi_table_function() -> None:

"""`with_session=True` is incompatible with FFI-exported table functions."""

class FakeFFITableFunction:

# Presence of this attribute is what marks a function as FFI-exported.

__datafusion_table_function__ = "stub"

fake = FakeFFITableFunction()

with pytest.raises(TypeError, match="FFI-exported table functions"):

udtf(fake, "fake_ffi", with_session=True)

with pytest.raises(TypeError, match="FFI-exported table functions"):

TableFunction("fake_ffi", fake, with_session=True)