cocoindex 0.2.23__tar.gz → 0.3.0__tar.gz

{cocoindex-0.2.23 → cocoindex-0.3.0}/Cargo.lock

@@ -1298,7 +1298,7 @@ dependencies = [
 
 [[package]]
 name = "cocoindex"
-version = "0.2.23"
+version = "0.3.0"
 dependencies = [
  "anyhow",
  "async-openai",
@@ -1363,6 +1363,7 @@ dependencies = [
  "time",
  "tokio",
  "tokio-stream",
+ "tokio-util",
  "tower 0.5.2",
  "tower-http",
  "tree-sitter",
@@ -5667,6 +5668,7 @@ dependencies = [
  "bytes",
  "futures-core",
  "futures-sink",
+ "futures-util",
  "pin-project-lite",
  "tokio",
 ]

{cocoindex-0.2.23 → cocoindex-0.3.0}/Cargo.toml

@@ -2,7 +2,7 @@
 name = "cocoindex"
 # Version used for local development is always higher than others to take precedence.
 # Will be overridden for specific release versions.
-version = "0.2.23"
+version = "0.3.0"
 edition = "2024"
 rust-version = "1.89"
 license = "Apache-2.0"
@@ -157,3 +157,4 @@ serde_path_to_error = "0.1.17"
 redis = { version = "0.31.0", features = ["tokio-comp", "connection-manager"] }
 expect-test = "1.5.0"
 encoding_rs = "0.8.35"
+tokio-util = { version = "0.7.16", features = ["rt"] }

{cocoindex-0.2.23 → cocoindex-0.3.0}/PKG-INFO

@@ -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 → cocoindex-0.3.0}/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>

{cocoindex-0.2.23 → cocoindex-0.3.0}/python/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-0.2.23 → cocoindex-0.3.0}/python/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-0.3.0/python/cocoindex/runtime.py

@@ -0,0 +1,74 @@
+"""
+This module provides a standalone execution runtime for executing coroutines in a thread-safe
+manner.
+"""
+
+import threading
+import asyncio
+import inspect
+import warnings
+
+from typing import Any, Callable, Awaitable, TypeVar, Coroutine
+
+T = TypeVar("T")
+
+
+class _ExecutionContext:
+    _lock: threading.Lock
+    _event_loop: asyncio.AbstractEventLoop | None = None
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+
+    @property
+    def event_loop(self) -> asyncio.AbstractEventLoop:
+        """Get the event loop for the cocoindex library."""
+        with self._lock:
+            if self._event_loop is None:
+                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."""
+        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()
+
+
+def to_async_call(call: Callable[..., Any]) -> Callable[..., Awaitable[Any]]:
+    if inspect.iscoroutinefunction(call):
+        return call
+    return lambda *args, **kwargs: asyncio.to_thread(lambda: call(*args, **kwargs))

{cocoindex-0.2.23 → cocoindex-0.3.0}/python/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 → cocoindex-0.3.0}/src/execution/live_updater.rs

@@ -10,7 +10,9 @@ use super::stats;
 use futures::future::try_join_all;
 use indicatif::ProgressBar;
 use sqlx::PgPool;
+use std::fmt::Write;
 use tokio::{sync::watch, task::JoinSet, time::MissedTickBehavior};
+use tokio_util::task::AbortOnDropHandle;
 
 pub struct FlowLiveUpdaterUpdates {
     pub active_sources: Vec<String>,
@@ -144,7 +146,6 @@ impl SourceUpdateTask {
             futs.push(
                 {
                     let change_stream_stats = change_stream_stats.clone();
-                    let pool = self.pool.clone();
                     let status_tx = self.status_tx.clone();
                     let operation_in_process_stats = self.operation_in_process_stats.clone();
                     async move {
@@ -229,7 +230,6 @@ impl SourceUpdateTask {
                                         Some(move || async move {
                                             SharedAckFn::ack(&shared_ack_fn).await
                                         }),
-                                        pool.clone(),
                                     ),
                                 );
                             }
@@ -252,7 +252,7 @@ impl SourceUpdateTask {
                         let curr_change_stream_stats = change_stream_stats.as_ref().clone();
                         let delta = curr_change_stream_stats.delta(&last_change_stream_stats);
                         if delta.has_any_change() {
-                            task.report_stats(&delta, "change stream");
+                            task.report_stats(&delta, "change stream", None);
                             last_change_stream_stats = curr_change_stream_stats;
                         }
                     }
@@ -266,7 +266,7 @@ impl SourceUpdateTask {
             async move {
                 let refresh_interval = import_op.refresh_options.refresh_interval;
 
-                task.update_with_pass_with_error_logging(
+                task.update_one_pass_with_error_logging(
                     source_indexing_context,
                     if refresh_interval.is_some() {
                         "initial interval update"
@@ -284,7 +284,7 @@ impl SourceUpdateTask {
                     loop {
                         interval.tick().await;
 
-                        task.update_with_pass_with_error_logging(
+                        task.update_one_pass_with_error_logging(
                             source_indexing_context,
                             "interval update",
                             super::source_indexer::UpdateOptions {
@@ -304,22 +304,30 @@ impl SourceUpdateTask {
         Ok(())
     }
 
-    fn report_stats(&self, stats: &stats::UpdateStats, update_title: &str) {
+    fn report_stats(
+        &self,
+        stats: &stats::UpdateStats,
+        update_title: &str,
+        start_time: Option<std::time::Instant>,
+    ) {
         self.source_update_stats.merge(stats);
+        let mut message = format!(
+            "{}.{} ({update_title}): {stats}",
+            self.flow.flow_instance.name,
+            self.import_op().name
+        );
+        if let Some(start_time) = start_time {
+            write!(
+                &mut message,
+                " [processing time: {:.3}s]",
+                start_time.elapsed().as_secs_f64()
+            )
+            .expect("Failed to write to message");
+        }
         if self.options.print_stats {
-            println!(
-                "{}.{} ({update_title}): {}",
-                self.flow.flow_instance.name,
-                self.import_op().name,
-                stats
-            );
+            println!("{message}");
         } else {
-            trace!(
-                "{}.{} ({update_title}): {}",
-                self.flow.flow_instance.name,
-                self.import_op().name,
-                stats
-            );
+            trace!("{message}");
         }
     }
 
@@ -329,6 +337,7 @@ impl SourceUpdateTask {
         update_title: &str,
         update_options: super::source_indexer::UpdateOptions,
     ) -> Result<()> {
+        let now = std::time::Instant::now();
         let update_stats = Arc::new(stats::UpdateStats::default());
 
         // Spawn periodic stats reporting task if print_stats is enabled
@@ -364,14 +373,17 @@ impl SourceUpdateTask {
                     }
                 }
             };
-            (Some(tokio::spawn(report_task)), Some(pb))
+            (
+                Some(AbortOnDropHandle::new(tokio::spawn(report_task))),
+                Some(pb),
+            )
         } else {
             (None, None)
         };
 
         // Run the actual update
         let update_result = source_indexing_context
-            .update(&self.pool, &update_stats, update_options)
+            .update(&update_stats, update_options)
             .await
             .with_context(|| {
                 format!(
@@ -401,11 +413,11 @@ impl SourceUpdateTask {
         }
 
         // Report final stats
-        self.report_stats(&update_stats, update_title);
+        self.report_stats(&update_stats, update_title, Some(now));
         Ok(())
     }
 
-    async fn update_with_pass_with_error_logging(
+    async fn update_one_pass_with_error_logging(
         &self,
         source_indexing_context: &Arc<SourceIndexingContext>,
         update_title: &str,