Merge branch 'cocoindex-io:main' into main

Author: Gohlub

docs/docs/custom_ops/custom_functions.mdx

@@ -145,6 +145,8 @@ Custom functions take the following additional parameters:
 * `batching: bool`: Whether the executor will consume requests in batch.
     See the [Batching](#batching) section below for details.
 
+* `max_batch_size: int | None`: The maximum batch size for the executor.
+
 * `behavior_version: int`: The version of the behavior of the function.
     When the version is changed, the function will be re-executed even if cache is enabled.
     It's required to be set if `cache` is `True`.
@@ -221,5 +223,25 @@ class ComputeSomethingExecutor:
         ...
 ```
 
+### Controlling Batch Size
+
+You can control the maximum batch size using the `max_batch_size` parameter. This is useful for:
+* Limiting memory usage when processing large batches
+* Reducing latency by flushing batches before they grow too large
+* Working with APIs that have request size limits
+
+```python
+@cocoindex.op.function(batching=True, max_batch_size=32)
+def compute_something(args: list[str]) -> list[str]:
+    ...
+```
+
+With `max_batch_size` set, a batch will be flushed when either:
+
+1. No ongoing batches are running
+2. The pending batch size reaches `max_batch_size`
+
+This ensures that requests don't wait indefinitely for a batch to fill up, while still allowing efficient batch processing.
+
 </TabItem>
 </Tabs>

docs/docs/sources/amazons3.md

@@ -131,6 +131,9 @@ The spec takes the following fields:
 
     :::
 
+*   `max_file_size` (`int`, optional): if provided, files exceeding this size in bytes will be treated as non-existent and skipped during processing.
+    This is useful to avoid processing large files that are not relevant to your use case, such as videos or backups.
+    If not specified, no size limit is applied.
 *   `sqs_queue_url` (`str`, optional): if provided, the source will receive change event notifications from Amazon S3 via this SQS queue.
 
     :::info

examples/code_embedding/main.py

@@ -16,12 +16,12 @@ def code_to_embedding(
     Embed the text using a SentenceTransformer model.
     """
     # You can also switch to Voyage embedding model:
-    #    return text.transform(
-    #        cocoindex.functions.EmbedText(
-    #            api_type=cocoindex.LlmApiType.VOYAGE,
-    #            model="voyage-code-3",
-    #        )
-    #    )
+    # return text.transform(
+    #     cocoindex.functions.EmbedText(
+    #         api_type=cocoindex.LlmApiType.GEMINI,
+    #         model="text-embedding-004",
+    #     )
+    # )
     return text.transform(
         cocoindex.functions.SentenceTransformerEmbed(
             model="sentence-transformers/all-MiniLM-L6-v2"

python/cocoindex/functions/colpali.py

@@ -125,6 +125,7 @@ class ColPaliEmbedImage(op.FunctionSpec):
     gpu=True,
     cache=True,
     batching=True,
+    max_batch_size=32,
     behavior_version=1,
 )
 class ColPaliEmbedImageExecutor:
@@ -204,6 +205,7 @@ class ColPaliEmbedQuery(op.FunctionSpec):
     cache=True,
     behavior_version=1,
     batching=True,
+    max_batch_size=32,
 )
 class ColPaliEmbedQueryExecutor:
     """Executor for ColVision query embedding (ColPali, ColQwen2, ColSmol, etc.)."""

python/cocoindex/functions/sbert.py

@@ -31,6 +31,7 @@ class SentenceTransformerEmbed(op.FunctionSpec):
     gpu=True,
     cache=True,
     batching=True,
+    max_batch_size=512,
     behavior_version=1,
     arg_relationship=(op.ArgRelationship.EMBEDDING_ORIGIN_TEXT, "text"),
 )

python/cocoindex/op.py

@@ -151,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.
+    - max_batch_size: The maximum batch size for the executor. Only valid if `batching` is True.
     - 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,
@@ -161,6 +162,7 @@ class OpArgs:
     gpu: bool = False
     cache: bool = False
     batching: bool = False
+    max_batch_size: int | None = None
     behavior_version: int | None = None
     arg_relationship: tuple[ArgRelationship, str] | None = None
 
@@ -389,11 +391,17 @@ def enable_cache(self) -> bool:
         def behavior_version(self) -> int | None:
             return op_args.behavior_version
 
+        def batching_options(self) -> dict[str, Any] | None:
+            if op_args.batching:
+                return {
+                    "max_batch_size": op_args.max_batch_size,
+                }
+            else:
+                return None
+
     if category == OpCategory.FUNCTION:
         _engine.register_function_factory(
-            op_kind,
-            _EngineFunctionExecutorFactory(spec_loader, _WrappedExecutor),
-            op_args.batching,
+            op_kind, _EngineFunctionExecutorFactory(spec_loader, _WrappedExecutor)
         )
     else:
         raise ValueError(f"Unsupported executor type {category}")

python/cocoindex/sources/_engine_builtin_specs.py

@@ -56,6 +56,7 @@ class AmazonS3(op.SourceSpec):
     binary: bool = False
     included_patterns: list[str] | None = None
     excluded_patterns: list[str] | None = None
+    max_file_size: int | None = None
     sqs_queue_url: str | None = None
     redis: RedisNotification | None = None
 

src/execution/source_indexer.rs

@@ -304,7 +304,10 @@ impl SourceIndexingContext {
                 rows_to_retry,
             }),
             setup_execution_ctx,
-            update_once_batcher: batching::Batcher::new(UpdateOnceRunner),
+            update_once_batcher: batching::Batcher::new(
+                UpdateOnceRunner,
+                batching::BatchingOptions::default(),
+            ),
         }))
     }
 

src/ops/factory_bases.rs

@@ -381,6 +381,8 @@ pub trait BatchedFunctionExecutor: Send + Sync + Sized + 'static {
     fn into_fn_executor(self) -> impl SimpleFunctionExecutor {
         BatchedFunctionExecutorWrapper::new(self)
     }
+
+    fn batching_options(&self) -> batching::BatchingOptions;
 }
 
 #[async_trait]
@@ -404,10 +406,11 @@ struct BatchedFunctionExecutorWrapper<E: BatchedFunctionExecutor> {
 
 impl<E: BatchedFunctionExecutor> BatchedFunctionExecutorWrapper<E> {
     fn new(executor: E) -> Self {
+        let batching_options = executor.batching_options();
         Self {
             enable_cache: executor.enable_cache(),
             behavior_version: executor.behavior_version(),
-            batcher: batching::Batcher::new(executor),
+            batcher: batching::Batcher::new(executor, batching_options),
         }
     }
 }

src/ops/functions/embed_text.rs

@@ -36,6 +36,14 @@ impl BatchedFunctionExecutor for Executor {
         true
     }
 
+    fn batching_options(&self) -> batching::BatchingOptions {
+        // A safe default for most embeddings providers.
+        // May tune it for specific providers later.
+        batching::BatchingOptions {
+            max_batch_size: Some(64),
+        }
+    }
+
     async fn evaluate_batch(&self, args: Vec<Vec<Value>>) -> Result<Vec<Value>> {
         let texts = args
             .iter()