test_udtf.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

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)


def test_pre_52_table_function_signature_reports_an_upgrade() -> None:
    """A getter that refuses the session is named, not left as a bare TypeError.

    The 52.0.0 signature change added the session argument. An out-of-date
    library still has the old one, and the original error stays reachable as
    ``__cause__``.
    """

    class PreSessionTableFunction:
        def __datafusion_table_function__(self):
            msg = "should never be called"
            raise AssertionError(msg)

    with pytest.raises(ImportError, match="__datafusion_table_function__") as excinfo:
        TableFunction("old_sig", PreSessionTableFunction(), None)
    assert isinstance(excinfo.value.__cause__, TypeError)


def test_type_error_inside_a_table_function_getter_propagates() -> None:
    """A correctly-signed getter's own TypeError must survive unchanged."""

    class RaisesTypeError:
        def __datafusion_table_function__(self, session):
            msg = "bad cast inside the getter"
            raise TypeError(msg)

    with pytest.raises(TypeError, match="bad cast inside the getter"):
        TableFunction("raises", RaisesTypeError(), None)