cocoindex 0.1.81__cp313-cp313-manylinux_2_28_aarch64.whl → 0.1.83__cp313-cp313-manylinux_2_28_aarch64.whl

cocoindex/op.py

@@ -17,14 +17,15 @@ from typing import (
 )
 
 from . import _engine  # type: ignore
+from .subprocess_exec import executor_stub
 from .convert import (
-    encode_engine_value,
+    make_engine_value_encoder,
     make_engine_value_decoder,
     make_engine_struct_decoder,
 )
 from .typing import (
     TypeAttr,
-    encode_enriched_type,
+    encode_enriched_type_info,
     resolve_forward_ref,
     analyze_type_info,
     AnalyzedAnyType,
@@ -85,11 +86,13 @@ class Executor(Protocol):
     op_category: OpCategory
 
 
-def _load_spec_from_engine(spec_cls: type, spec: dict[str, Any]) -> Any:
+def _load_spec_from_engine(
+    spec_loader: Callable[..., Any], spec: dict[str, Any]
+) -> Any:
     """
     Load a spec from the engine.
     """
-    return spec_cls(**spec)
+    return spec_loader(**spec)
 
 
 def _get_required_method(cls: type, name: str) -> Callable[..., Any]:
@@ -101,26 +104,23 @@ def _get_required_method(cls: type, name: str) -> Callable[..., Any]:
     return method
 
 
-class _FunctionExecutorFactory:
-    _spec_cls: type
+class _EngineFunctionExecutorFactory:
+    _spec_loader: Callable[..., Any]
     _executor_cls: type
 
-    def __init__(self, spec_cls: type, executor_cls: type):
-        self._spec_cls = spec_cls
+    def __init__(self, spec_loader: Callable[..., Any], executor_cls: type):
+        self._spec_loader = spec_loader
         self._executor_cls = executor_cls
 
     def __call__(
         self, spec: dict[str, Any], *args: Any, **kwargs: Any
     ) -> tuple[dict[str, Any], Executor]:
-        spec = _load_spec_from_engine(self._spec_cls, spec)
+        spec = _load_spec_from_engine(self._spec_loader, spec)
         executor = self._executor_cls(spec)
         result_type = executor.analyze_schema(*args, **kwargs)
         return (result_type, executor)
 
 
-_gpu_dispatch_lock = asyncio.Lock()
-
-
 _COCOINDEX_ATTR_PREFIX = "cocoindex.io/"
 
 
@@ -166,30 +166,32 @@ def _register_op_factory(
     category: OpCategory,
     expected_args: list[tuple[str, inspect.Parameter]],
     expected_return: Any,
-    executor_cls: type,
-    spec_cls: type,
+    executor_factory: Any,
+    spec_loader: Callable[..., Any],
+    op_kind: str,
     op_args: OpArgs,
-) -> type:
+) -> None:
     """
     Register an op factory.
     """
 
-    class _Fallback:
-        def enable_cache(self) -> bool:
-            return op_args.cache
-
-        def behavior_version(self) -> int | None:
-            return op_args.behavior_version
-
-    class _WrappedClass(executor_cls, _Fallback):  # type: ignore[misc]
+    class _WrappedExecutor:
+        _executor: Any
         _args_info: list[_ArgInfo]
         _kwargs_info: dict[str, _ArgInfo]
-        _acall: Callable[..., Awaitable[Any]]
+        _result_encoder: Callable[[Any], Any]
+        _acall: Callable[..., Awaitable[Any]] | None = None
 
         def __init__(self, spec: Any) -> None:
-            super().__init__()
-            self.spec = spec
-            self._acall = _to_async_call(super().__call__)
+            executor: Any
+
+            if op_args.gpu:
+                executor = executor_stub(executor_factory, spec)
+            else:
+                executor = executor_factory()
+                executor.spec = spec
+
+            self._executor = executor
 
         def analyze_schema(
             self, *args: _engine.OpArgSchema, **kwargs: _engine.OpArgSchema
@@ -293,17 +295,21 @@ def _register_op_factory(
             if len(missing_args) > 0:
                 raise ValueError(f"Missing arguments: {', '.join(missing_args)}")
 
-            base_analyze_method = getattr(self, "analyze", None)
+            base_analyze_method = getattr(self._executor, "analyze", None)
             if base_analyze_method is not None:
-                result = base_analyze_method(*args, **kwargs)
+                result_type = base_analyze_method()
             else:
-                result = expected_return
+                result_type = expected_return
             if len(attributes) > 0:
-                result = Annotated[result, *attributes]
+                result_type = Annotated[result_type, *attributes]
 
-            encoded_type = encode_enriched_type(result)
+            analyzed_result_type_info = analyze_type_info(result_type)
+            encoded_type = encode_enriched_type_info(analyzed_result_type_info)
             if potentially_missing_required_arg:
                 encoded_type["nullable"] = True
+
+            self._result_encoder = make_engine_value_encoder(analyzed_result_type_info)
+
             return encoded_type
 
         async def prepare(self) -> None:
@@ -311,9 +317,10 @@ def _register_op_factory(
             Prepare for execution.
             It's executed after `analyze` and before any `__call__` execution.
             """
-            prepare_method = getattr(super(), "prepare", None)
+            prepare_method = getattr(self._executor, "prepare", None)
             if prepare_method is not None:
                 await _to_async_call(prepare_method)()
+            self._acall = _to_async_call(self._executor.__call__)
 
         async def __call__(self, *args: Any, **kwargs: Any) -> Any:
             decoded_args = []
@@ -333,33 +340,23 @@ def _register_op_factory(
                     return None
                 decoded_kwargs[kwarg_name] = kwarg_info.decoder(arg)
 
-            if op_args.gpu:
-                # For GPU executions, data-level parallelism is applied, so we don't want to
-                # execute different tasks in parallel.
-                # Besides, multiprocessing is more appropriate for pytorch.
-                # For now, we use a lock to ensure only one task is executed at a time.
-                # TODO: Implement multi-processing dispatching.
-                async with _gpu_dispatch_lock:
-                    output = await self._acall(*decoded_args, **decoded_kwargs)
-            else:
-                output = await self._acall(*decoded_args, **decoded_kwargs)
-            return encode_engine_value(output)
+            assert self._acall is not None
+            output = await self._acall(*decoded_args, **decoded_kwargs)
+            return self._result_encoder(output)
 
-    _WrappedClass.__name__ = executor_cls.__name__
-    _WrappedClass.__doc__ = executor_cls.__doc__
-    _WrappedClass.__module__ = executor_cls.__module__
-    _WrappedClass.__qualname__ = executor_cls.__qualname__
-    _WrappedClass.__wrapped__ = executor_cls
+        def enable_cache(self) -> bool:
+            return op_args.cache
+
+        def behavior_version(self) -> int | None:
+            return op_args.behavior_version
 
     if category == OpCategory.FUNCTION:
         _engine.register_function_factory(
-            spec_cls.__name__, _FunctionExecutorFactory(spec_cls, _WrappedClass)
+            op_kind, _EngineFunctionExecutorFactory(spec_loader, _WrappedExecutor)
         )
     else:
         raise ValueError(f"Unsupported executor type {category}")
 
-    return _WrappedClass
-
 
 def executor_class(**args: Any) -> Callable[[type], type]:
     """
@@ -377,52 +374,53 @@ def executor_class(**args: Any) -> Callable[[type], type]:
             raise TypeError("Expect a `spec` field with type hint")
         spec_cls = resolve_forward_ref(type_hints["spec"])
         sig = inspect.signature(cls.__call__)
-        return _register_op_factory(
+        _register_op_factory(
             category=spec_cls._op_category,
             expected_args=list(sig.parameters.items())[1:],  # First argument is `self`
             expected_return=sig.return_annotation,
-            executor_cls=cls,
-            spec_cls=spec_cls,
+            executor_factory=cls,
+            spec_loader=spec_cls,
+            op_kind=spec_cls.__name__,
             op_args=op_args,
         )
+        return cls
 
     return _inner
 
 
-def function(**args: Any) -> Callable[[Callable[..., Any]], FunctionSpec]:
+class EmptyFunctionSpec(FunctionSpec):
+    pass
+
+
+class _SimpleFunctionExecutor:
+    spec: Callable[..., Any]
+
+    def prepare(self) -> None:
+        self.__call__ = staticmethod(self.spec)
+
+
+def function(**args: Any) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
     """
     Decorate a function to provide a function for an op.
     """
     op_args = OpArgs(**args)
 
-    def _inner(fn: Callable[..., Any]) -> FunctionSpec:
+    def _inner(fn: Callable[..., Any]) -> Callable[..., Any]:
         # Convert snake case to camel case.
-        op_name = "".join(word.capitalize() for word in fn.__name__.split("_"))
+        op_kind = "".join(word.capitalize() for word in fn.__name__.split("_"))
         sig = inspect.signature(fn)
-
-        class _Executor:
-            def __call__(self, *args: Any, **kwargs: Any) -> Any:
-                return fn(*args, **kwargs)
-
-        class _Spec(FunctionSpec):
-            def __call__(self, *args: Any, **kwargs: Any) -> Any:
-                return fn(*args, **kwargs)
-
-        _Spec.__name__ = op_name
-        _Spec.__doc__ = fn.__doc__
-        _Spec.__module__ = fn.__module__
-        _Spec.__qualname__ = fn.__qualname__
-
+        fn.__cocoindex_op_kind__ = op_kind  # type: ignore
         _register_op_factory(
             category=OpCategory.FUNCTION,
             expected_args=list(sig.parameters.items()),
             expected_return=sig.return_annotation,
-            executor_cls=_Executor,
-            spec_cls=_Spec,
+            executor_factory=_SimpleFunctionExecutor,
+            spec_loader=lambda: fn,
+            op_kind=op_kind,
             op_args=op_args,
         )
 
-        return _Spec()
+        return fn
 
     return _inner
 

cocoindex/sources.py

@@ -2,6 +2,7 @@
 
 from . import op
 from .auth_registry import TransientAuthEntryReference
+from .setting import DatabaseConnectionSpec
 import datetime
 
 
@@ -67,3 +68,22 @@ class AzureBlob(op.SourceSpec):
 
     sas_token: TransientAuthEntryReference[str] | None = None
     account_access_key: TransientAuthEntryReference[str] | None = None
+
+
+class Postgres(op.SourceSpec):
+    """Import data from a PostgreSQL table."""
+
+    _op_category = op.OpCategory.SOURCE
+
+    # Table name to read from (required)
+    table_name: str
+
+    # Database connection reference (optional - uses default if not provided)
+    database: TransientAuthEntryReference[DatabaseConnectionSpec] | None = None
+
+    # Optional: specific columns to include (if None, includes all columns)
+    included_columns: list[str] | None = None
+
+    # Optional: column name to use for ordinal tracking (for incremental updates)
+    # Should be a timestamp, serial, or other incrementing column
+    ordinal_column: str | None = None

cocoindex/subprocess_exec.py

@@ -0,0 +1,241 @@
+"""
+Lightweight subprocess-backed executor stub.
+
+- Uses a single global ProcessPoolExecutor (max_workers=1), created lazily.
+- In the subprocess, maintains a registry of executor instances keyed by
+  (executor_factory, pickled spec) to enable reuse.
+- Caches analyze() and prepare() results per key to avoid repeated calls
+  even if key collision happens.
+"""
+
+from __future__ import annotations
+
+from concurrent.futures import ProcessPoolExecutor
+from concurrent.futures.process import BrokenProcessPool
+from dataclasses import dataclass, field
+from typing import Any, Callable
+import pickle
+import threading
+import asyncio
+import os
+import time
+from .user_app_loader import load_user_app
+from .runtime import execution_context
+import logging
+
+WATCHDOG_INTERVAL_SECONDS = 10.0
+
+# ---------------------------------------------
+# Main process: single, lazily-created pool
+# ---------------------------------------------
+_pool_lock = threading.Lock()
+_pool: ProcessPoolExecutor | None = None
+_user_apps: list[str] = []
+_logger = logging.getLogger(__name__)
+
+
+def _get_pool() -> ProcessPoolExecutor:
+    global _pool
+    with _pool_lock:
+        if _pool is None:
+            # Single worker process as requested
+            _pool = ProcessPoolExecutor(
+                max_workers=1,
+                initializer=_subprocess_init,
+                initargs=(_user_apps, os.getpid()),
+            )
+        return _pool
+
+
+def add_user_app(app_target: str) -> None:
+    with _pool_lock:
+        _user_apps.append(app_target)
+
+
+def _restart_pool(old_pool: ProcessPoolExecutor | None = None) -> None:
+    """Safely restart the global ProcessPoolExecutor.
+
+    Thread-safe via `_pool_lock`. Shuts down the old pool and re-creates a new
+    one with the same initializer/args.
+    """
+    global _pool
+    with _pool_lock:
+        # If another thread already swapped the pool, skip restart
+        if old_pool is not None and _pool is not old_pool:
+            return
+        _logger.error("Detected dead subprocess pool; restarting and retrying.")
+        prev_pool = _pool
+        _pool = ProcessPoolExecutor(
+            max_workers=1,
+            initializer=_subprocess_init,
+            initargs=(_user_apps, os.getpid()),
+        )
+        if prev_pool is not None:
+            # Best-effort shutdown of previous pool; letting exceptions bubble up
+            # is acceptable here and signals irrecoverable executor state.
+            prev_pool.shutdown(cancel_futures=True)
+
+
+async def _submit_with_restart(fn: Callable[..., Any], *args: Any) -> Any:
+    """Submit and await work, restarting the subprocess until it succeeds.
+
+    Retries on BrokenProcessPool or pool-shutdown RuntimeError; re-raises other
+    exceptions.
+    """
+    while True:
+        pool = _get_pool()
+        try:
+            fut = pool.submit(fn, *args)
+            return await asyncio.wrap_future(fut)
+        except BrokenProcessPool:
+            _restart_pool(old_pool=pool)
+            # loop and retry
+
+
+# ---------------------------------------------
+# Subprocess: executor registry and helpers
+# ---------------------------------------------
+
+
+def _start_parent_watchdog(
+    parent_pid: int, interval_seconds: float = WATCHDOG_INTERVAL_SECONDS
+) -> None:
+    """Terminate this process if the parent process exits or PPID changes.
+
+    This runs in a background daemon thread so it never blocks pool work.
+    """
+
+    def _watch() -> None:
+        while True:
+            # If PPID changed (parent died and we were reparented), exit.
+            if os.getppid() != parent_pid:
+                os._exit(1)
+
+            # Best-effort liveness probe in case PPID was reused.
+            try:
+                os.kill(parent_pid, 0)
+            except OSError:
+                os._exit(1)
+
+            time.sleep(interval_seconds)
+
+    threading.Thread(target=_watch, name="parent-watchdog", daemon=True).start()
+
+
+def _subprocess_init(user_apps: list[str], parent_pid: int) -> None:
+    _start_parent_watchdog(parent_pid)
+    for app_target in user_apps:
+        load_user_app(app_target)
+
+
+class _OnceResult:
+    _result: Any = None
+    _done: bool = False
+
+    def run_once(self, method: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
+        if self._done:
+            return self._result
+        self._result = _call_method(method, *args, **kwargs)
+        self._done = True
+        return self._result
+
+
+@dataclass
+class _ExecutorEntry:
+    executor: Any
+    prepare: _OnceResult = field(default_factory=_OnceResult)
+    analyze: _OnceResult = field(default_factory=_OnceResult)
+    ready_to_call: bool = False
+
+
+_SUBPROC_EXECUTORS: dict[bytes, _ExecutorEntry] = {}
+
+
+def _call_method(method: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
+    """Run an awaitable/coroutine to completion synchronously, otherwise return as-is."""
+    if asyncio.iscoroutinefunction(method):
+        return asyncio.run(method(*args, **kwargs))
+    else:
+        return method(*args, **kwargs)
+
+
+def _get_or_create_entry(key_bytes: bytes) -> _ExecutorEntry:
+    entry = _SUBPROC_EXECUTORS.get(key_bytes)
+    if entry is None:
+        executor_factory, spec = pickle.loads(key_bytes)
+        inst = executor_factory()
+        inst.spec = spec
+        entry = _ExecutorEntry(executor=inst)
+        _SUBPROC_EXECUTORS[key_bytes] = entry
+    return entry
+
+
+def _sp_analyze(key_bytes: bytes) -> Any:
+    entry = _get_or_create_entry(key_bytes)
+    return entry.analyze.run_once(entry.executor.analyze)
+
+
+def _sp_prepare(key_bytes: bytes) -> Any:
+    entry = _get_or_create_entry(key_bytes)
+    return entry.prepare.run_once(entry.executor.prepare)
+
+
+def _sp_call(key_bytes: bytes, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
+    entry = _get_or_create_entry(key_bytes)
+    # There's a chance that the subprocess crashes and restarts in the middle.
+    # So we want to always make sure the executor is ready before each call.
+    if not entry.ready_to_call:
+        if analyze_fn := getattr(entry.executor, "analyze", None):
+            entry.analyze.run_once(analyze_fn)
+        if prepare_fn := getattr(entry.executor, "prepare", None):
+            entry.prepare.run_once(prepare_fn)
+        entry.ready_to_call = True
+    return _call_method(entry.executor.__call__, *args, **kwargs)
+
+
+# ---------------------------------------------
+# Public stub
+# ---------------------------------------------
+
+
+class _ExecutorStub:
+    _pool: ProcessPoolExecutor
+    _key_bytes: bytes
+
+    def __init__(self, executor_factory: type[Any], spec: Any) -> None:
+        self._pool = _get_pool()
+        self._key_bytes = pickle.dumps(
+            (executor_factory, spec), protocol=pickle.HIGHEST_PROTOCOL
+        )
+
+        # Conditionally expose analyze if underlying class has it
+        if hasattr(executor_factory, "analyze"):
+            # Bind as attribute so getattr(..., "analyze", None) works upstream
+            def analyze() -> Any:
+                return execution_context.run(
+                    _submit_with_restart(_sp_analyze, self._key_bytes)
+                )
+
+            # Attach method
+            setattr(self, "analyze", analyze)
+
+        if hasattr(executor_factory, "prepare"):
+
+            async def prepare() -> Any:
+                return await _submit_with_restart(_sp_prepare, self._key_bytes)
+
+            setattr(self, "prepare", prepare)
+
+    async def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        return await _submit_with_restart(_sp_call, self._key_bytes, args, kwargs)
+
+
+def executor_stub(executor_factory: type[Any], spec: Any) -> Any:
+    """
+    Create a subprocess-backed stub for the given executor class/spec.
+
+    - Lazily initializes a singleton ProcessPoolExecutor (max_workers=1).
+    - Returns a stub object exposing async __call__ and async prepare; analyze is
+      exposed if present on the original class.
+    """
+    return _ExecutorStub(executor_factory, spec)