vgi-python 0.34.1__py3-none-any.whl → 0.36.0__py3-none-any.whl

vgi/_test_fixtures/table/filters.py

@@ -1003,7 +1003,10 @@ class FilterEchoPartitionedFunction(TableFunctionGenerator[_FilterEchoPartitione
             end_idx = min(start_idx + chunk, params.args.count)
             work_items.append(struct.pack(">QQ", start_idx, end_idx))
         params.storage.queue_push(work_items)
-        return GlobalInitResponse()
+        # No more readers than work items. Left at the default (unbounded) the
+        # client opens one stream per DuckDB thread -- 48 on a 48-core host for at
+        # most MAX_PARTITIONS items, every extra one an init plus an empty drain.
+        return GlobalInitResponse(max_workers=max(1, len(work_items)))
 
     @classmethod
     def initial_state(cls, params: ProcessParams[_FilterEchoPartitionedArgs]) -> _FilterEchoPartitionedState:

vgi/_test_fixtures/table/sequence.py

@@ -516,7 +516,10 @@ class PartitionedSequenceFunction(
 
         # Always enqueue (even if empty) to register the invocation
         params.storage.queue_push(work_items)
-        return GlobalInitResponse()
+        # No more readers than work items. Left at the default (unbounded) the
+        # client opens one stream per DuckDB thread -- 48 on a 48-core host for at
+        # most MAX_PARTITIONS items, every extra one an init plus an empty drain.
+        return GlobalInitResponse(max_workers=max(1, len(work_items)))
 
     @classmethod
     def initial_state(cls, params: ProcessParams[PartitionedSequenceFunctionArguments]) -> PartitionedSequenceState:

vgi/_test_fixtures/table/splits.py

@@ -35,6 +35,7 @@ from vgi_rpc.rpc import OutputCollector
 
 from vgi.arguments import Arg
 from vgi.cache_control import CacheControl
+from vgi.filter_v2 import replaying_accepted_state
 from vgi.metadata import FunctionExample, PartitionKind
 from vgi.protocol import PlanResponse, ScanSplit, VgiOutputCollector
 from vgi.schema_utils import partition_field, schema
@@ -1055,11 +1056,15 @@ class SplitDynamicFilterFunction(TableFunctionGenerator[SplitSequenceArgs, Split
             init = params.init_call
             merged = None
             if init is not None and init.pushdown_filters is not None:
-                merged = cls.pushdown_filters(
-                    init.pushdown_filters,
-                    join_keys=init.join_keys,
-                    output_schema=init.output_schema,
-                )
+                # The framework validated these at init; re-binding every
+                # predicate through the evaluator on every tick is what made
+                # this fixture cost ~10 ms a batch.
+                with replaying_accepted_state():
+                    merged = cls.pushdown_filters(
+                        init.pushdown_filters,
+                        join_keys=init.join_keys,
+                        output_schema=init.output_schema,
+                    )
             if merged is None:
                 merged = params.current_pushdown_filters
             filter_str = _render_filters_canonical(merged)

vgi/_test_fixtures/table_in_out.py

@@ -138,9 +138,9 @@ class SubstreamPartialSumFunction(TableInOutFunction[SingleTableArguments, Subst
     their finalize outputs, so the caller re-aggregates with an outer
     ``SELECT sum(...)`` to get the global total — correct no matter how the rows
     were partitioned across substreams. Each substream's ``finish()`` reads only
-    its OWN worker's accumulated state (keyed by the substream's execution_id;
-    ``params.substream_id`` is the stable client-owned key available for workers
-    that manage cross-backend state themselves). This is the per-substream
+    its OWN execution's accumulated state — the framework keeps one row per
+    ``params.substream_id`` in storage scoped to that execution, so it is found
+    even when a finalize lands on a different backend. This is the per-substream
     finalize contract A4 enables — it is NOT a global cross-substream combine
     (that is a ``TableBufferingFunction``; see ``SumAllColumnsSimpleDistributed``).
 
@@ -191,8 +191,9 @@ class SubstreamPartialSumFunction(TableInOutFunction[SingleTableArguments, Subst
         params: ProcessParams[SingleTableArguments],
         states: list[SubstreamPartialSumState],
     ) -> list[pa.RecordBatch]:
-        # `states` are THIS substream's accumulated states (one per worker pid that
-        # handled this substream's batches); their sum is this substream's partial.
+        # `states` are this execution's accumulated states, one per substream that
+        # saw input (the Python client fans one execution across several); their
+        # sum is this finalize's partial.
         total = sum(st.total for st in states)
         name = params.output_schema.names[0]
         return [pa.RecordBatch.from_pydict({name: [total]}, schema=params.output_schema)]

vgi/auth.py

@@ -18,8 +18,9 @@ JWT auth (requires ``vgi[oauth]``):
 
 from __future__ import annotations
 
-import contextlib
+import importlib
 from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
 
 from vgi_rpc.rpc import AuthContext, CallContext
 
@@ -43,34 +44,89 @@ __all__ = [
     "TokenResolver",
 ]
 
-# HTTP auth helpers — available when vgi[http] is installed.
-with contextlib.suppress(ImportError):
-    from vgi_rpc.http import (  # noqa: F401
-        AuthUnavailableError,
-        OAuthResourceMetadata,
-        bearer_authenticate,
-        bearer_authenticate_static,
-        chain_authenticate,
-        parse_client_id,
-        parse_client_secret,
-        parse_device_code_client_id,
-        parse_device_code_client_secret,
+# HTTP auth helpers (``vgi[http]``) and JWT auth (``vgi[oauth]``), resolved on
+# first access rather than at import. ``vgi/__init__`` imports this module, and
+# ``vgi_rpc.http`` brings in the HTTP server and client stack (falcon, httpx2,
+# cryptography, joserfc): importing it here put about 180 ms on every
+# ``import vgi``, and so on the start of every subprocess worker, which never
+# uses it. ``from vgi.auth import bearer_authenticate`` works as before; when
+# the extra is missing it raises ImportError, as it did.
+_LAZY_EXPORTS = {
+    "AuthUnavailableError": "vgi_rpc.http",
+    "OAuthResourceMetadata": "vgi_rpc.http",
+    "bearer_authenticate": "vgi_rpc.http",
+    "bearer_authenticate_static": "vgi_rpc.http",
+    "chain_authenticate": "vgi_rpc.http",
+    "parse_client_id": "vgi_rpc.http",
+    "parse_client_secret": "vgi_rpc.http",
+    "parse_device_code_client_id": "vgi_rpc.http",
+    "parse_device_code_client_secret": "vgi_rpc.http",
+    "jwt_authenticate": "vgi_rpc.http._oauth_jwt",
+}
+
+if TYPE_CHECKING:
+    # ``X as X``: explicit re-exports, so type checkers see these as part of
+    # ``vgi.auth`` exactly as they were before they became lazy.
+    from vgi_rpc.http import (
+        AuthUnavailableError as AuthUnavailableError,
+    )
+    from vgi_rpc.http import (
+        OAuthResourceMetadata as OAuthResourceMetadata,
+    )
+    from vgi_rpc.http import (
+        bearer_authenticate as bearer_authenticate,
+    )
+    from vgi_rpc.http import (
+        bearer_authenticate_static as bearer_authenticate_static,
+    )
+    from vgi_rpc.http import (
+        chain_authenticate as chain_authenticate,
+    )
+    from vgi_rpc.http import (
+        parse_client_id as parse_client_id,
+    )
+    from vgi_rpc.http import (
+        parse_client_secret as parse_client_secret,
+    )
+    from vgi_rpc.http import (
+        parse_device_code_client_id as parse_device_code_client_id,
     )
+    from vgi_rpc.http import (
+        parse_device_code_client_secret as parse_device_code_client_secret,
+    )
+    from vgi_rpc.http._oauth_jwt import jwt_authenticate as jwt_authenticate
+
+
+def __getattr__(name: str) -> Any:
+    """Import an HTTP or JWT auth helper on first access (PEP 562).
+
+    Args:
+        name: The attribute being looked up.
+
+    Returns:
+        The helper, which is then cached in this module's namespace.
+
+    Raises:
+        AttributeError: If ``name`` is not a helper this module re-exports, or
+            its extra is not installed.
+    """
+    module_name = _LAZY_EXPORTS.get(name)
+    if module_name is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    try:
+        value = getattr(importlib.import_module(module_name), name)
+    except ImportError as exc:
+        raise AttributeError(
+            f"module {__name__!r} has no attribute {name!r}: it needs {module_name}, which failed to import ({exc})"
+        ) from exc
+    globals()[name] = value
+    return value
+
+
+def __dir__() -> list[str]:
+    """List this module's names, the lazily imported helpers included.
 
-    __all__ += [
-        "AuthUnavailableError",
-        "OAuthResourceMetadata",
-        "bearer_authenticate",
-        "bearer_authenticate_static",
-        "chain_authenticate",
-        "parse_client_id",
-        "parse_client_secret",
-        "parse_device_code_client_id",
-        "parse_device_code_client_secret",
-    ]
-
-# JWT auth — available when vgi[oauth] is installed (requires authlib).
-with contextlib.suppress(ImportError):
-    from vgi_rpc.http._oauth_jwt import jwt_authenticate  # noqa: F401
-
-    __all__ += ["jwt_authenticate"]
+    Returns:
+        The module's attribute names.
+    """
+    return sorted({*globals(), *_LAZY_EXPORTS})

vgi/catalog/catalog_interface.py

@@ -2698,6 +2698,8 @@ class ReadOnlyCatalogInterface(CatalogInterface):
     _function_registry: "dict[tuple[SchemaKey, str], list[type]] | None" = None
     _macro_registry: "dict[tuple[SchemaKey, str], Macro] | None" = None
     _index_registry: "dict[tuple[SchemaKey, str], Index] | None" = None
+    # Function listings by (schema, listing type), filled in by ``_function_infos``.
+    _function_info_cache: "dict[tuple[SchemaKey, SchemaObjectType], tuple[FunctionInfo, ...]] | None" = None
     # Lazy registry build is one-time but the fixture HTTP server is
     # multi-threaded and shares one catalog instance, so concurrent
     # first-requests can race the build. Serialize it under a lock and flip
@@ -2738,6 +2740,7 @@ class ReadOnlyCatalogInterface(CatalogInterface):
         self._function_registry = {}
         self._macro_registry = {}
         self._index_registry = {}
+        self._function_info_cache = {}
 
         def _register_table(schema_key: SchemaKey, table: "Table") -> None:
             key = (schema_key, table.name.lower())
@@ -3320,29 +3323,46 @@ class ReadOnlyCatalogInterface(CatalogInterface):
                 if sn == path_key and macro.macro_type == target_macro_type:
                     results.append(macro.to_macro_info(schema_path))
         else:
-            # SCALAR_FUNCTION or TABLE_FUNCTION
-            for (sn, _), func_classes in self._function_registry.items():
-                if sn != path_key:
-                    continue
-                for func_cls in func_classes:
-                    func_info = self._function_to_info(func_cls, schema_path)
-                    # Filter by function type
-                    if type_enum == SchemaObjectType.SCALAR_FUNCTION and func_info.function_type != FunctionType.SCALAR:
-                        continue
-                    if type_enum == SchemaObjectType.TABLE_FUNCTION and func_info.function_type not in (
-                        FunctionType.TABLE,
-                        FunctionType.TABLE_BUFFERING,
-                    ):
-                        continue
-                    if (
-                        type_enum == SchemaObjectType.AGGREGATE_FUNCTION
-                        and func_info.function_type != FunctionType.AGGREGATE
-                    ):
-                        continue
-                    results.append(func_info)
+            # SCALAR_FUNCTION, TABLE_FUNCTION or AGGREGATE_FUNCTION
+            results.extend(self._function_infos(path_key, schema_path, type_enum))
 
         return results
 
+    def _function_infos(
+        self, path_key: SchemaKey, schema_path: SchemaPath, type_enum: SchemaObjectType
+    ) -> tuple[FunctionInfo, ...]:
+        """Return the schema's functions of one listing type, built on the first request.
+
+        A listing depends only on the registered classes' static metadata, so
+        every request for it returns the same `FunctionInfo` instances. Rebuilding
+        it on each request re-derived every function in the schema to keep the
+        requested type's. It also made the worker encode each one again, at
+        about 0.8 ms per function, because ``FunctionsResponse.from_infos``
+        stores an instance's encoding on the instance.
+        The instances are shared between requests, so they must not be mutated.
+        """
+        assert self._function_registry is not None
+        assert self._function_info_cache is not None
+        key = (path_key, type_enum)
+        cached = self._function_info_cache.get(key)
+        if cached is not None:
+            return cached
+        wanted = {
+            SchemaObjectType.SCALAR_FUNCTION: (FunctionType.SCALAR,),
+            SchemaObjectType.TABLE_FUNCTION: (FunctionType.TABLE, FunctionType.TABLE_BUFFERING),
+            SchemaObjectType.AGGREGATE_FUNCTION: (FunctionType.AGGREGATE,),
+        }.get(type_enum, ())
+        infos: list[FunctionInfo] = []
+        for (sn, _), func_classes in self._function_registry.items():
+            if sn != path_key:
+                continue
+            for func_cls in func_classes:
+                func_info = self._function_to_info(func_cls, schema_path)
+                if func_info.function_type in wanted:
+                    infos.append(func_info)
+        # Concurrent first requests may both build it; keep whichever landed first.
+        return self._function_info_cache.setdefault(key, tuple(infos))
+
     def copy_from_formats(
         self,
         *,