table functions continued

Author: rustyconover

tests/table/generator/conftest.py

@@ -0,0 +1,89 @@
+"""Shared fixtures for table generator function tests."""
+
+from collections.abc import Callable, Generator
+from typing import Any, Literal
+
+import pyarrow as pa
+import pytest
+
+from vgi.client import Client, ClientError
+from vgi.function import Arguments
+from vgi.log import Message
+from vgi.table_function import TableFunctionGenerator
+from vgi.testing import FunctionTestClientError, TableFunctionTestClient
+
+# Type alias for the runner function
+TableFunctionRunner = Callable[
+    [type[TableFunctionGenerator], tuple[Any, ...]],
+    tuple[list[pa.RecordBatch], list[Message]],
+]
+
+# Type alias for the test mode
+TestMode = Literal["in_process", "client"]
+
+# Type alias for the fixture return type
+RunnerWithMode = tuple[TableFunctionRunner, TestMode]
+
+
+def run_in_process(
+    func_class: type[TableFunctionGenerator],
+    args: tuple[Any, ...],
+) -> tuple[list[pa.RecordBatch], list[Message]]:
+    """Run a table function in-process using TableFunctionTestClient."""
+    with TableFunctionTestClient(func_class) as client:
+        outputs = list(
+            client.table_function(
+                arguments=Arguments(positional=tuple(pa.scalar(a) for a in args))
+            )
+        )
+        return outputs, client.logs
+
+
+def run_via_client(
+    func_class: type[TableFunctionGenerator],
+    args: tuple[Any, ...],
+) -> tuple[list[pa.RecordBatch], list[Message]]:
+    """Run a table function via subprocess using Client.table_function.
+
+    Uses max_workers=1 to ensure consistent behavior with in-process mode.
+    For multi-worker tests, use Client directly with explicit max_workers.
+    """
+    meta = func_class.get_metadata()
+    function_name = meta.name
+
+    with Client("vgi-example-worker", max_workers=1) as client:
+        outputs = list(
+            client.table_function(
+                function_name=function_name,
+                arguments=Arguments(positional=tuple(pa.scalar(a) for a in args)),
+            )
+        )
+        # Note: logs are not captured via Client (they go to stderr)
+        return outputs, []
+
+
+@pytest.fixture(params=["in_process", "client"])
+def run_table_function_mode(
+    request: pytest.FixtureRequest,
+) -> Generator[RunnerWithMode, None, None]:
+    """Fixture that provides both in-process and client-based runners.
+
+    Tests using this fixture will run twice: once in-process and once via Client.
+    Returns a tuple of (runner_function, mode_name) for conditional assertions.
+    """
+    mode: TestMode = request.param
+    if mode == "in_process":
+        yield run_in_process, mode
+    else:
+        yield run_via_client, mode
+
+
+@pytest.fixture
+def run_function() -> TableFunctionRunner:
+    """Fixture for in-process-only tests (e.g., log capture tests)."""
+    return run_in_process
+
+
+# Re-export error types for convenience
+InProcessError = FunctionTestClientError
+SubprocessError = ClientError

tests/table/generator/test_partitioned_function.py

@@ -0,0 +1,158 @@
+"""Tests for the PartitionedRangeFunction with multi-worker support."""
+
+import pyarrow as pa
+
+from vgi.client import Client
+from vgi.examples.table import PartitionedRangeFunction
+from vgi.function import Arguments
+from vgi.testing import TableFunctionTestClient
+
+from .conftest import RunnerWithMode
+
+
+class TestPartitionedRangeFunctionInProcess:
+    """In-process tests for the partitioned_range function."""
+
+    def test_generates_full_range_single_worker(self) -> None:
+        """Single worker should generate the complete range."""
+        with TableFunctionTestClient(PartitionedRangeFunction) as client:
+            outputs = list(
+                client.table_function(
+                    arguments=Arguments(positional=(pa.scalar(10),))
+                )
+            )
+
+        table = pa.Table.from_batches(outputs)
+        assert table.num_rows == 10
+
+        values = sorted(table.column("value").to_pylist())
+        assert values == list(range(10))
+
+    def test_metadata(self) -> None:
+        """Partitioned range function should have correct metadata."""
+        meta = PartitionedRangeFunction.get_metadata()
+        assert meta.name == "partitioned_range"
+        # Should not have max_workers limit (parallelizable)
+        assert meta.max_workers is None
+
+    def test_zero_count(self) -> None:
+        """Partitioned range with count=0 should produce no output."""
+        with TableFunctionTestClient(PartitionedRangeFunction) as client:
+            outputs = list(
+                client.table_function(
+                    arguments=Arguments(positional=(pa.scalar(0),))
+                )
+            )
+
+        assert len(outputs) == 0
+
+
+class TestPartitionedRangeFunctionBothModes:
+    """Tests that run both in-process and via Client subprocess."""
+
+    def test_generates_correct_count(
+        self, run_table_function_mode: RunnerWithMode
+    ) -> None:
+        """Partitioned range should generate exactly the requested number of rows."""
+        runner, mode = run_table_function_mode
+        outputs, logs = runner(PartitionedRangeFunction, (100,))
+
+        table = pa.Table.from_batches(outputs)
+        assert table.num_rows == 100
+
+    def test_values_are_sequential(
+        self, run_table_function_mode: RunnerWithMode
+    ) -> None:
+        """Partitioned range should produce all values in range."""
+        runner, mode = run_table_function_mode
+        outputs, logs = runner(PartitionedRangeFunction, (50,))
+
+        table = pa.Table.from_batches(outputs)
+        values = sorted(table.column("value").to_pylist())
+        assert values == list(range(50))
+
+
+class TestPartitionedRangeFunctionMultiWorker:
+    """Tests for multi-worker partitioned execution via Client."""
+
+    def test_two_workers_produce_complete_range(self) -> None:
+        """Two workers should together produce the complete range."""
+        with Client("vgi-example-worker", max_workers=2) as client:
+            outputs = list(
+                client.table_function(
+                    function_name="partitioned_range",
+                    arguments=Arguments(positional=(pa.scalar(20),)),
+                )
+            )
+
+        table = pa.Table.from_batches(outputs)
+        assert table.num_rows == 20
+
+        values = sorted(table.column("value").to_pylist())
+        assert values == list(range(20))
+
+    def test_three_workers_produce_complete_range(self) -> None:
+        """Three workers should together produce the complete range."""
+        with Client("vgi-example-worker", max_workers=3) as client:
+            outputs = list(
+                client.table_function(
+                    function_name="partitioned_range",
+                    arguments=Arguments(positional=(pa.scalar(30),)),
+                )
+            )
+
+        table = pa.Table.from_batches(outputs)
+        assert table.num_rows == 30
+
+        values = sorted(table.column("value").to_pylist())
+        assert values == list(range(30))
+
+    def test_workers_produce_large_range(self) -> None:
+        """Multiple workers should handle large ranges."""
+        with Client("vgi-example-worker", max_workers=4) as client:
+            outputs = list(
+                client.table_function(
+                    function_name="partitioned_range",
+                    arguments=Arguments(positional=(pa.scalar(10000),)),
+                )
+            )
+
+        table = pa.Table.from_batches(outputs)
+        assert table.num_rows == 10000
+
+        values = sorted(table.column("value").to_pylist())
+        assert values == list(range(10000))
+
+    def test_uneven_distribution(self) -> None:
+        """Workers should handle ranges that don't divide evenly."""
+        # 7 items with 3 workers: worker 0 gets [0,3,6], worker 1 gets [1,4],
+        # worker 2 gets [2,5]
+        with Client("vgi-example-worker", max_workers=3) as client:
+            outputs = list(
+                client.table_function(
+                    function_name="partitioned_range",
+                    arguments=Arguments(positional=(pa.scalar(7),)),
+                )
+            )
+
+        table = pa.Table.from_batches(outputs)
+        assert table.num_rows == 7
+
+        values = sorted(table.column("value").to_pylist())
+        assert values == list(range(7))
+
+    def test_single_worker_fallback(self) -> None:
+        """max_workers=1 should work like single worker mode."""
+        with Client("vgi-example-worker", max_workers=1) as client:
+            outputs = list(
+                client.table_function(
+                    function_name="partitioned_range",
+                    arguments=Arguments(positional=(pa.scalar(15),)),
+                )
+            )
+
+        table = pa.Table.from_batches(outputs)
+        assert table.num_rows == 15
+
+        values = sorted(table.column("value").to_pylist())
+        assert values == list(range(15))