cocoindex 0.2.23__cp311-abi3-macosx_10_12_x86_64.whl → 0.3.0__cp311-abi3-macosx_10_12_x86_64.whl

cocoindex/functions/sbert.py

@@ -31,6 +31,7 @@ class SentenceTransformerEmbed(op.FunctionSpec):
 @op.executor_class(
     gpu=True,
     cache=True,
+    batching=True,
     behavior_version=1,
     arg_relationship=(op.ArgRelationship.EMBEDDING_ORIGIN_TEXT, "text"),
 )
@@ -57,7 +58,9 @@ class SentenceTransformerEmbedExecutor:
         dim = self._model.get_sentence_embedding_dimension()
         return Vector[np.float32, Literal[dim]]  # type: ignore
 
-    def __call__(self, text: str) -> NDArray[np.float32]:
+    def __call__(self, text: list[str]) -> list[NDArray[np.float32]]:
         assert self._model is not None
-        result: NDArray[np.float32] = self._model.encode(text, convert_to_numpy=True)
-        return result
+        results: list[NDArray[np.float32]] = self._model.encode(
+            text, convert_to_numpy=True
+        )
+        return results

cocoindex/op.py

@@ -32,6 +32,7 @@ from .engine_value import (
 )
 from .typing import (
     KEY_FIELD_NAME,
+    AnalyzedListType,
     AnalyzedTypeInfo,
     StructSchema,
     StructType,
@@ -45,6 +46,7 @@ from .typing import (
     EnrichedValueType,
     decode_engine_field_schemas,
     FieldSchema,
+    ValueType,
 )
 from .runtime import to_async_call
 from .index import IndexOptions
@@ -149,6 +151,7 @@ class OpArgs:
     """
     - gpu: Whether the executor will be executed on GPU.
     - cache: Whether the executor will be cached.
+    - batching: Whether the executor will be batched.
     - behavior_version: The behavior version of the executor. Cache will be invalidated if it
       changes. Must be provided if `cache` is True.
     - arg_relationship: It specifies the relationship between an input argument and the output,
@@ -158,6 +161,7 @@ class OpArgs:
 
     gpu: bool = False
     cache: bool = False
+    batching: bool = False
     behavior_version: int | None = None
     arg_relationship: tuple[ArgRelationship, str] | None = None
 
@@ -168,6 +172,16 @@ class _ArgInfo:
     is_required: bool
 
 
+def _make_batched_engine_value_decoder(
+    field_path: list[str], src_type: ValueType, dst_type_info: AnalyzedTypeInfo
+) -> Callable[[Any], Any]:
+    if not isinstance(dst_type_info.variant, AnalyzedListType):
+        raise ValueError("Expected arguments for batching function to be a list type")
+    elem_type_info = analyze_type_info(dst_type_info.variant.elem_type)
+    base_decoder = make_engine_value_decoder(field_path, src_type, elem_type_info)
+    return lambda value: [base_decoder(v) for v in value]
+
+
 def _register_op_factory(
     category: OpCategory,
     expected_args: list[tuple[str, inspect.Parameter]],
@@ -181,6 +195,10 @@ def _register_op_factory(
     Register an op factory.
     """
 
+    if op_args.batching:
+        if len(expected_args) != 1:
+            raise ValueError("Batching is only supported for single argument functions")
+
     class _WrappedExecutor:
         _executor: Any
         _args_info: list[_ArgInfo]
@@ -208,7 +226,7 @@ def _register_op_factory(
             """
             self._args_info = []
             self._kwargs_info = {}
-            attributes = []
+            attributes = {}
             potentially_missing_required_arg = False
 
             def process_arg(
@@ -220,14 +238,17 @@ def _register_op_factory(
                 if op_args.arg_relationship is not None:
                     related_attr, related_arg_name = op_args.arg_relationship
                     if related_arg_name == arg_name:
-                        attributes.append(
-                            TypeAttr(related_attr.value, actual_arg.analyzed_value)
-                        )
+                        attributes[related_attr.value] = actual_arg.analyzed_value
                 type_info = analyze_type_info(arg_param.annotation)
                 enriched = EnrichedValueType.decode(actual_arg.value_type)
-                decoder = make_engine_value_decoder(
-                    [arg_name], enriched.type, type_info
-                )
+                if op_args.batching:
+                    decoder = _make_batched_engine_value_decoder(
+                        [arg_name], enriched.type, type_info
+                    )
+                else:
+                    decoder = make_engine_value_decoder(
+                        [arg_name], enriched.type, type_info
+                    )
                 is_required = not type_info.nullable
                 if is_required and actual_arg.value_type.get("nullable", False):
                     potentially_missing_required_arg = True
@@ -302,20 +323,32 @@ def _register_op_factory(
             if len(missing_args) > 0:
                 raise ValueError(f"Missing arguments: {', '.join(missing_args)}")
 
+            analyzed_expected_return_type = analyze_type_info(expected_return)
+            self._result_encoder = make_engine_value_encoder(
+                analyzed_expected_return_type
+            )
+
             base_analyze_method = getattr(self._executor, "analyze", None)
             if base_analyze_method is not None:
-                result_type = base_analyze_method()
+                analyzed_result_type = analyze_type_info(base_analyze_method())
             else:
-                result_type = expected_return
+                if op_args.batching:
+                    if not isinstance(
+                        analyzed_expected_return_type.variant, AnalyzedListType
+                    ):
+                        raise ValueError(
+                            "Expected return type for batching function to be a list type"
+                        )
+                    analyzed_result_type = analyze_type_info(
+                        analyzed_expected_return_type.variant.elem_type
+                    )
+                else:
+                    analyzed_result_type = analyzed_expected_return_type
             if len(attributes) > 0:
-                result_type = Annotated[result_type, *attributes]
-
-            analyzed_result_type_info = analyze_type_info(result_type)
-            encoded_type = encode_enriched_type_info(analyzed_result_type_info)
+                analyzed_result_type.attrs = attributes
             if potentially_missing_required_arg:
-                encoded_type["nullable"] = True
-
-            self._result_encoder = make_engine_value_encoder(analyzed_result_type_info)
+                analyzed_result_type.nullable = True
+            encoded_type = encode_enriched_type_info(analyzed_result_type)
 
             return encoded_type
 
@@ -359,7 +392,9 @@ def _register_op_factory(
 
     if category == OpCategory.FUNCTION:
         _engine.register_function_factory(
-            op_kind, _EngineFunctionExecutorFactory(spec_loader, _WrappedExecutor)
+            op_kind,
+            _EngineFunctionExecutorFactory(spec_loader, _WrappedExecutor),
+            op_args.batching,
         )
     else:
         raise ValueError(f"Unsupported executor type {category}")

cocoindex/runtime.py

@@ -6,8 +6,9 @@ manner.
 import threading
 import asyncio
 import inspect
-from typing import Any, Callable, Coroutine, TypeVar, Awaitable
+import warnings
 
+from typing import Any, Callable, Awaitable, TypeVar, Coroutine
 
 T = TypeVar("T")
 
@@ -24,15 +25,44 @@ class _ExecutionContext:
         """Get the event loop for the cocoindex library."""
         with self._lock:
             if self._event_loop is None:
-                self._event_loop = asyncio.new_event_loop()
-                threading.Thread(
-                    target=self._event_loop.run_forever, daemon=True
-                ).start()
+                loop = asyncio.new_event_loop()
+                self._event_loop = loop
+
+                def _runner(l: asyncio.AbstractEventLoop) -> None:
+                    asyncio.set_event_loop(l)
+                    l.run_forever()
+
+                threading.Thread(target=_runner, args=(loop,), daemon=True).start()
             return self._event_loop
 
     def run(self, coro: Coroutine[Any, Any, T]) -> T:
         """Run a coroutine in the event loop, blocking until it finishes. Return its result."""
-        return asyncio.run_coroutine_threadsafe(coro, self.event_loop).result()
+        try:
+            running_loop = asyncio.get_running_loop()
+        except RuntimeError:
+            running_loop = None
+
+        loop = self.event_loop
+
+        if running_loop is not None:
+            if running_loop is loop:
+                raise RuntimeError(
+                    "CocoIndex sync API was called from inside CocoIndex's async context. "
+                    "Use the async variant of this method instead."
+                )
+            warnings.warn(
+                "CocoIndex sync API was called inside an existing event loop. "
+                "This may block other tasks. Prefer the async method.",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+
+        fut = asyncio.run_coroutine_threadsafe(coro, loop)
+        try:
+            return fut.result()
+        except KeyboardInterrupt:
+            fut.cancel()
+            raise
 
 
 execution_context = _ExecutionContext()

cocoindex/subprocess_exec.py

@@ -19,7 +19,6 @@ import threading
 import asyncio
 import os
 import time
-import atexit
 from .user_app_loader import load_user_app
 from .runtime import execution_context
 import logging
@@ -32,39 +31,14 @@ WATCHDOG_INTERVAL_SECONDS = 10.0
 # ---------------------------------------------
 _pool_lock = threading.Lock()
 _pool: ProcessPoolExecutor | None = None
-_pool_cleanup_registered = False
 _user_apps: list[str] = []
 _logger = logging.getLogger(__name__)
 
 
-def shutdown_pool_at_exit() -> None:
-    """Best-effort shutdown of the global ProcessPoolExecutor on interpreter exit."""
-    global _pool, _pool_cleanup_registered  # pylint: disable=global-statement
-    with _pool_lock:
-        if _pool is not None:
-            try:
-                _pool.shutdown(wait=True, cancel_futures=True)
-            except Exception as e:
-                _logger.error(
-                    "Error during ProcessPoolExecutor shutdown at exit: %s",
-                    e,
-                    exc_info=True,
-                )
-            finally:
-                _pool = None
-                _pool_cleanup_registered = False
-
-
 def _get_pool() -> ProcessPoolExecutor:
-    global _pool, _pool_cleanup_registered  # pylint: disable=global-statement
+    global _pool  # pylint: disable=global-statement
     with _pool_lock:
         if _pool is None:
-            if not _pool_cleanup_registered:
-                # Register the shutdown at exit at creation time (rather than at import time)
-                # to make sure it's executed earlier in the shutdown sequence.
-                atexit.register(shutdown_pool_at_exit)
-                _pool_cleanup_registered = True
-
             # Single worker process as requested
             _pool = ProcessPoolExecutor(
                 max_workers=1,

{cocoindex-0.2.23.dist-info → cocoindex-0.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cocoindex
-Version: 0.2.23
+Version: 0.3.0
 Classifier: Development Status :: 3 - Alpha
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent

{cocoindex-0.2.23.dist-info → cocoindex-0.3.0.dist-info}/RECORD

@@ -1,9 +1,9 @@
-cocoindex-0.2.23.dist-info/METADATA,sha256=ngGBMCTqt5Apsew33CJ7osgLGBwAdSBr8E3etjUEFB0,14194
-cocoindex-0.2.23.dist-info/WHEEL,sha256=eDlp9unULyyDxD2Zd14qZwSC_Y-kO5nuRBdlMsoCXEY,105
-cocoindex-0.2.23.dist-info/entry_points.txt,sha256=_NretjYVzBdNTn7dK-zgwr7YfG2afz1u1uSE-5bZXF8,46
-cocoindex-0.2.23.dist-info/licenses/THIRD_PARTY_NOTICES.html,sha256=JkjceWLwTbf-svSD5D0wLkWQTF2nCHB_Ev3THD6IWOQ,750832
+cocoindex-0.3.0.dist-info/METADATA,sha256=vMA1fqUWkDkxXgNT76l8pOBEqgt-Cu8XEpC4BPCbzp8,14193
+cocoindex-0.3.0.dist-info/WHEEL,sha256=eDlp9unULyyDxD2Zd14qZwSC_Y-kO5nuRBdlMsoCXEY,105
+cocoindex-0.3.0.dist-info/entry_points.txt,sha256=_NretjYVzBdNTn7dK-zgwr7YfG2afz1u1uSE-5bZXF8,46
+cocoindex-0.3.0.dist-info/licenses/THIRD_PARTY_NOTICES.html,sha256=VRDb6qOsN808v5fkXRUNfGUqSvD_OJWjO6hD3uWfKFg,750831
 cocoindex/__init__.py,sha256=6qZWVkK4WZ01BIAg3CPh_bRRdA6Clk4d4Q6OnZ2jFa4,2630
-cocoindex/_engine.abi3.so,sha256=E3eoWnJh9lz9OdOtJ5vmcRjNqCoq8Jf_9jMCFRp1tBk,70055200
+cocoindex/_engine.abi3.so,sha256=CHhZSvV-OpSq20F8vg7-kY6ibRQ-ahlCZ2FQw52c77E,70186592
 cocoindex/auth_registry.py,sha256=g-uLDWLYW5NMbYe7q4Y-sU5dSyrlJXBEciyWtAiP9KE,1340
 cocoindex/cli.py,sha256=k7bl8RTUZoNNxTlQMr-Y3-9-rTNt8z1v7rJWqsajYC8,24792
 cocoindex/engine_object.py,sha256=5YTuWoR3WILhyt3PW-d9es3MAas_xD6tZZqvipN-sjg,10050
@@ -12,19 +12,19 @@ cocoindex/flow.py,sha256=xDz3rOo4RhbboknvC-KnbWq8RBykEO0YsjGSBfXqIEg,40076
 cocoindex/functions/__init__.py,sha256=V2IF4h-Cqq4OD_GN3Oqdry-FArORyRCKmqJ7g5UlJr8,1021
 cocoindex/functions/_engine_builtin_specs.py,sha256=WpCGrjUfJBa8xZP5JiEmA8kLu7fp9Rcs7ynpuJmvSGg,1786
 cocoindex/functions/colpali.py,sha256=oACyG3qG2dquyCJ6bT7FkMkua5rXDLSxnOHcgoz9waU,8865
-cocoindex/functions/sbert.py,sha256=1z5OJT-blXT6tVN5vEvEzvYAzOnzs1RCnu1UbCUP6wM,2162
+cocoindex/functions/sbert.py,sha256=o_DS1ZAqpNc4u1Yrm9DO5LxfMFrlH_hfb0MWobJrs_k,2223
 cocoindex/index.py,sha256=tz5ilvmOp0BtroGehCQDqWK_pIX9m6ghkhcxsDVU8WE,982
 cocoindex/lib.py,sha256=spfdU4IbzdffHyGdrQPIw_qGo9aX0OAAboqsjj8bTiQ,2290
 cocoindex/llm.py,sha256=8ZdJhOmhdb2xEcCxk6rDpnj6hlhCyFBmJdhCNMqAOP4,875
-cocoindex/op.py,sha256=TO-ETk3qXgnNS51NlWuLrOw_TfQ2mw83-_iswqULcQI,36095
+cocoindex/op.py,sha256=FtYHn3RPBWSBAAaFbmp-FhLbs0qrMMd3lZdheM-udUo,37672
 cocoindex/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cocoindex/query_handler.py,sha256=X-SQT71LHiOOXn6-TJlQcGodJk-iT8p_1TcIMvRLBRI,1344
-cocoindex/runtime.py,sha256=4NxcltaDZvA3RR3Pnt6gH_f99jcWSyMH_1Xi5BjbtwY,1342
+cocoindex/runtime.py,sha256=Qig--_3Uwvcw-MuuEdBg-P92IYB6So_nRocD2jnEi3g,2298
 cocoindex/setting.py,sha256=1Dx8ktjwf-8BiXrbsmfn5Mzudb2SQYqFdRnSNGVKaLk,4960
 cocoindex/setup.py,sha256=7uIHKN4FOCuoidPXcKyGTrkqpkl9luL49-6UcnMxYzw,3068
 cocoindex/sources/__init__.py,sha256=Yu9VHNaGlOEE3jpqfIseswsg25Le3HzwDr6XJAn22Ns,78
 cocoindex/sources/_engine_builtin_specs.py,sha256=MDsNP1y0vlUCGTAmA4jj2X8AMn1MVATbZzSi66ggvkA,3598
-cocoindex/subprocess_exec.py,sha256=r1xO84uek4VP4I6i87JMwsH5xFm3vKW0ABvgn0jskt4,10088
+cocoindex/subprocess_exec.py,sha256=41PyvMaZ9q7B-F0YTUg1xdNVM_6_HdFErDbxDJMUFPg,9040
 cocoindex/targets/__init__.py,sha256=HQG7I4U0xQhHiYctiUvwEBLxT2727oHP3xwrqotjmhk,78
 cocoindex/targets/_engine_builtin_specs.py,sha256=glXUN5bj11Jxky1VPvmGnWnMHXTQWEh08INcbldo3F4,3375
 cocoindex/targets/lancedb.py,sha256=1nzCre5p-fvKkmLOTvfpiLTfnhF3qMLqTvsTwNuGwVU,15749
@@ -39,4 +39,4 @@ cocoindex/typing.py,sha256=qQj5uM6XAKHzRJ2BIEs7X-xeOXVcM9p_xz5SVqPVvS8,23914
 cocoindex/user_app_loader.py,sha256=bc3Af-gYRxJ9GpObtpjegZY855oQBCv5FGkrkWV2yGY,1873
 cocoindex/utils.py,sha256=hUhX-XV6XGCtJSEIpBOuDv6VvqImwPlgBxztBTw7u0U,598
 cocoindex/validation.py,sha256=PZnJoby4sLbsmPv9fOjOQXuefjfZ7gmtsiTGU8SH-tc,3090
-cocoindex-0.2.23.dist-info/RECORD,,
+cocoindex-0.3.0.dist-info/RECORD,,

{cocoindex-0.2.23.dist-info → cocoindex-0.3.0.dist-info}/licenses/THIRD_PARTY_NOTICES.html

@@ -2846,7 +2846,7 @@ Software.
                 <h3 id="Apache-2.0">Apache License 2.0</h3>
                 <h4>Used by:</h4>
                 <ul class="license-used-by">
-                    <li><a href=" https://crates.io/crates/cocoindex ">cocoindex 0.2.23</a></li>
+                    <li><a href=" https://crates.io/crates/cocoindex ">cocoindex 0.3.0</a></li>
                     <li><a href=" https://github.com/awesomized/crc-fast-rust ">crc-fast 1.3.0</a></li>
                     <li><a href=" https://github.com/qdrant/rust-client ">qdrant-client 1.15.0</a></li>
                 </ul>