cocoindex 0.2.23__tar.gz → 0.3.1__tar.gz

{cocoindex-0.2.23 → cocoindex-0.3.1}/.pre-commit-config.yaml

@@ -76,3 +76,4 @@ repos:
         types: [python]
         pass_filenames: false
         always_run: false
+        args: ["python"]

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

@@ -1298,7 +1298,7 @@ dependencies = [
 
 [[package]]
 name = "cocoindex"
-version = "0.2.23"
+version = "0.3.1"
 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.1}/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.1"
 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.1}/PKG-INFO

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

{cocoindex-0.2.23 → cocoindex-0.3.1}/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.1</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.1}/docs/docs/custom_ops/custom_functions.mdx

@@ -8,8 +8,8 @@ import TabItem from '@theme/TabItem';
 
 A custom function can be defined in one of the following ways:
 
-*   A standalone function. It's simpler and doesn't allow additional configurations and setup logic.
-*   A function spec and an executor. It's more powerful, allows additional configurations and setup logic.
+* A standalone function. It's simpler and doesn't allow additional configurations and setup logic.
+* A function spec and an executor. It's more powerful, allows additional configurations and setup logic.
 
 ## Option 1: By a standalone function
 
@@ -31,9 +31,9 @@ def compute_something(arg1: str, arg2: int | None = None) -> str:
 
 Notes:
 
-*   The `cocoindex.op.function()` function decorator also takes optional parameters.
+* The `cocoindex.op.function()` function decorator also takes optional parameters.
     See [Parameters for custom functions](#parameters-for-custom-functions) for details.
-*   Types of arguments and the return value must be annotated, so that CocoIndex will have information about data types of the operation's output fields.
+* Types of arguments and the return value must be annotated, so that CocoIndex will have information about data types of the operation's output fields.
     See [Data Types](/docs/core/data_types) for supported types.
 
 </TabItem>
@@ -43,12 +43,11 @@ Notes:
 
 The cocoindex repository contains the following examples of custom functions defined in this way:
 
-*   In the [code_embedding](https://github.com/cocoindex-io/cocoindex/blob/main/examples/code_embedding/main.py) example,
+* In the [code_embedding](https://github.com/cocoindex-io/cocoindex/blob/main/examples/code_embedding/main.py) example,
     `extract_extension` is a custom function to extract the extension of a file name.
-*   In the [manuals_llm_extraction](https://github.com/cocoindex-io/cocoindex/blob/main/examples/manuals_llm_extraction/main.py) example,
+* In the [manuals_llm_extraction](https://github.com/cocoindex-io/cocoindex/blob/main/examples/manuals_llm_extraction/main.py) example,
     `summarize_manuals` is a custom function to summarize structured information of a manual page.
 
-
 ## Option 2: By a function spec and an executor
 
 This is more advanced and flexible way to define a custom function.
@@ -75,30 +74,29 @@ class ComputeSomething(cocoindex.op.FunctionSpec):
 ```
 
 Notes:
-*   All fields of the spec must have a type serializable / deserializable by the `json` module.
-*   All subclasses of `FunctionSpec` can be instantiated similar to a dataclass, i.e. `ClassName(param1=value1, param2=value2, ...)`.
+
+* All fields of the spec must have a type serializable / deserializable by the `json` module.
+* All subclasses of `FunctionSpec` can be instantiated similar to a dataclass, i.e. `ClassName(param1=value1, param2=value2, ...)`.
 
 </TabItem>
 </Tabs>
 
-
 ### Function Executor
 
 A function executor defines behavior of a function. It's instantiated for each operation that uses this function.
 
 The function executor is responsible for:
 
-*   *Prepare* for the function execution, based on the spec.
+* *Prepare* for the function execution, based on the spec.
     It happens once and only once before execution.
     e.g. if the function calls a machine learning model, the model name can be a parameter as a field of the spec, and we may load the model in this phase.
-*   *Run* the function, for each specific input arguments. This happens multiple times, for each specific row of data.
+* *Run* the function, for each specific input arguments. This happens multiple times, for each specific row of data.
 
 <Tabs>
 <TabItem value="python" label="Python" default>
 
 A function executor is defined as a class decorated by `@cocoindex.op.executor_class()`.
 
-
 ```python
 @cocoindex.op.executor_class(...)
 class ComputeSomethingExecutor:
@@ -114,12 +112,12 @@ class ComputeSomethingExecutor:
 
 Notes:
 
-*   The `cocoindex.op.executor_class()` class decorator also takes optional parameters.
+* The `cocoindex.op.executor_class()` class decorator also takes optional parameters.
     See [Parameters for custom functions](#parameters-for-custom-functions) for details.
 
-*   A `spec` field must be present in the class, and must be annotated with the spec class name.
-*   The `prepare()` method is optional. It's executed once and only once before any `__call__` execution, to prepare the function execution.
-*   The `__call__()` method is required. It's executed for each specific rows of data.
+* A `spec` field must be present in the class, and must be annotated with the spec class name.
+* The `prepare()` method is optional. It's executed once and only once before any `__call__` execution, to prepare the function execution.
+* The `__call__()` method is required. It's executed for each specific rows of data.
     Types of arugments and the return value must be decorated, so that CocoIndex will have information about data types of the operation's output fields.
     See [Data Types](/docs/core/data_types) for supported types.
 
@@ -130,34 +128,37 @@ Notes:
 
 The cocoindex repository contains the following examples of custom functions defined in this way:
 
-*   In the [pdf_embedding](https://github.com/cocoindex-io/cocoindex/blob/main/examples/pdf_embedding/main.py) example, we define a custom function `PdfToMarkdown`
-*   The `SentenceTransformerEmbed` function shipped with the CocoIndex Python package is defined by Python SDK.
+* In the [pdf_embedding](https://github.com/cocoindex-io/cocoindex/blob/main/examples/pdf_embedding/main.py) example, we define a custom function `PdfToMarkdown`
+* The `SentenceTransformerEmbed` function shipped with the CocoIndex Python package is defined by Python SDK.
     Search for [`SentenceTransformerEmbedExecutor`](https://github.com/search?q=repo%3Acocoindex-io%2Fcocoindex+lang%3Apython+SentenceTransformerEmbedExecutor&type=code) to see the code.
 
 ## Parameters for custom functions
 
 Custom functions take the following additional parameters:
 
-*   `gpu: bool`: Whether the executor will use GPU. It will affect the way the function is scheduled.
+* `gpu: bool`: Whether the executor will use GPU. It will affect the way the function is scheduled.
 
-*   `cache: bool`: Whether the executor will enable cache for this function.
+* `cache: bool`: Whether the executor will enable cache for this function.
      When `True`, the executor will cache the result of the function for reuse during reprocessing.
      We recommend to set this to `True` for any function that is computationally intensive.
 
-*   `behavior_version: int`: The version of the behavior of the function.
+* `batching: bool`: Whether the executor will consume requests in batch.
+    See the [Batching](#batching) section below for details.
+
+* `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`.
 
-*   `arg_relationship: tuple[ArgRelationship, str]`: It specifies the relationship between an input argument and the output,
+* `arg_relationship: tuple[ArgRelationship, str]`: It specifies the relationship between an input argument and the output,
     e.g. `(ArgRelationship.CHUNKS_BASE_TEXT, "content")` means the output is chunks for the text represented by the
     input argument with name `content`.
     This provides metadata for tools, e.g. CocoInsight.
     Currently the following attributes are supported:
 
-    *   `ArgRelationship.CHUNKS_BASE_TEXT`:
+  * `ArgRelationship.CHUNKS_BASE_TEXT`:
         The output is chunks for the text represented by the input argument. In this case, the output is expected to be a *Table*, whose each row represents a text chunk, and the first column has type *Range*, representing the range of the text chunk.
-    *   `ArgRelationship.EMBEDDING_ORIGIN_TEXT`: The output is embedding vector for the text represented by the input argument. The output is expected to be a *Vector*.
-    *   `ArgRelationship.RECTS_BASE_IMAGE`: The output is rectangles for the image represented by the input argument. The output is expected to be a *Table*, whose each row represents a rectangle, and the first column has type *Struct*, with fields `min_x`, `min_y`, `max_x`, `max_y` to represent the coordinates of the rectangle.
+  * `ArgRelationship.EMBEDDING_ORIGIN_TEXT`: The output is embedding vector for the text represented by the input argument. The output is expected to be a *Vector*.
+  * `ArgRelationship.RECTS_BASE_IMAGE`: The output is rectangles for the image represented by the input argument. The output is expected to be a *Table*, whose each row represents a rectangle, and the first column has type *Struct*, with fields `min_x`, `min_y`, `max_x`, `max_y` to represent the coordinates of the rectangle.
 
 For example:
 
@@ -187,3 +188,38 @@ class ComputeSomethingExecutor:
 
 </TabItem>
 </Tabs>
+
+## Batching
+
+Batching allows a function executor to process multiple function calls in batch.
+Sometimes batching is more efficient than processing them one by one, e.g. running inference on GPU, calling remote APIs with quota limits, etc.
+
+Batching can be enabled by setting the `batching` parameter to `True` in custom function parameters.
+Once it's set to `True`, type of the argument and return value must be a `list`.
+Currently we only support batching functions taking a single argument.
+
+<Tabs>
+<TabItem value="python" label="Python" default>
+
+For example, for a CocoIndex custom function taking `str` as argument and returning `str` as result, it can be defined as:
+
+```python
+@cocoindex.op.function(batching=True)
+def compute_something(args: list[str]) -> list[str]:
+    ...
+```
+
+or for such a function defined by an executor:
+
+```python
+@cocoindex.op.executor_class(batching=True)
+class ComputeSomethingExecutor:
+    spec: ComputeSomething
+    ...
+
+    def __call__(self, args: list[str]) -> list[str]:
+        ...
+```
+
+</TabItem>
+</Tabs>

{cocoindex-0.2.23 → cocoindex-0.3.1}/docs/docs/examples/examples/01_docs_to_knowledge_graph.md

@@ -373,4 +373,4 @@ You can open it at [http://localhost:7474](http://localhost:7474), and run the f
 MATCH p=()-->() RETURN p
 ```
 
-![Neo4j Browser](/img/examples/docs_to_knowledge_graph/neo4j_browser.png)
+![Neo4j Browser](/img/examples/docs_to_knowledge_graph/neo4j_browser.png)

{cocoindex-0.2.23 → cocoindex-0.3.1}/docs/docs/sources/index.md

@@ -17,6 +17,6 @@ In CocoIndex, a source is the data origin you import from (e.g., files, database
 | [Postgres](/docs/sources/postgres)       | Relational database (Postgres)           |
 
 Related:
-- [Life cycle of a indexing flow](/docs/core/basics#life-cycle-of-an-indexing-flow) 
-- [Live Update Tutorial](/docs/tutorials/live_updates) 
+- [Life cycle of a indexing flow](/docs/core/basics#life-cycle-of-an-indexing-flow)
+- [Live Update Tutorial](/docs/tutorials/live_updates)
 for change capture mechanisms.

{cocoindex-0.2.23 → cocoindex-0.3.1}/docs/docs/targets/index.md

@@ -334,6 +334,3 @@ You can find end-to-end examples fitting into any of supported property graphs i
 *   <ExampleButton href="https://github.com/cocoindex-io/cocoindex/tree/main/examples/docs_to_knowledge_graph" text="Docs to Knowledge Graph" margin="0 0 16px 0" />
 
 *   <ExampleButton href="https://github.com/cocoindex-io/cocoindex/tree/main/examples/product_recommendation" text="Product Recommendation" margin="0 0 16px 0" />
-
-
-

{cocoindex-0.2.23 → cocoindex-0.3.1}/docs/docs/targets/kuzu.md

@@ -13,7 +13,7 @@ Exports data to a [Kuzu](https://kuzu.com/) graph database.
 
 ## Get Started
 
-Read [Property Graph Targets](./index.md#property-graph-targets) for more information to get started on how it works in CocoIndex. 
+Read [Property Graph Targets](./index.md#property-graph-targets) for more information to get started on how it works in CocoIndex.
 
 ## Spec
 
@@ -59,4 +59,4 @@ You can then access the explorer at [http://localhost:8124](http://localhost:812
   href="https://github.com/cocoindex-io/cocoindex/tree/main/examples/docs_to_knowledge_graph"
   text="Docs to Knowledge Graph"
   margin="16px 0 24px 0"
-/>
+/>

{cocoindex-0.2.23 → cocoindex-0.3.1}/docs/docs/targets/neo4j.md

@@ -11,7 +11,7 @@ import { ExampleButton } from '../../src/components/GitHubButton';
 
 
 ## Get Started
-Read [Property Graph Targets](./index.md#property-graph-targets) for more information to get started on how it works in CocoIndex. 
+Read [Property Graph Targets](./index.md#property-graph-targets) for more information to get started on how it works in CocoIndex.
 
 
 ## Spec
@@ -59,4 +59,4 @@ If you are building multiple CocoIndex flows from different projects to neo4j, w
 
 This way, you can clean up the data for each flow independently.
 
-In case you need to clean up the data in the same database, you can do it manually by running `cocoindex drop <APP_TARGET>` from the project you want to clean up.
+In case you need to clean up the data in the same database, you can do it manually by running `cocoindex drop <APP_TARGET>` from the project you want to clean up.

{cocoindex-0.2.23 → cocoindex-0.3.1}/examples/product_recommendation/README.md

@@ -8,7 +8,7 @@ Please drop [CocoIndex on Github](https://github.com/cocoindex-io/cocoindex) a s
 
 
 ## Prerequisite
-*   [Install Postgres](https://cocoindex.io/docs/getting_started/installation#-install-postgres) 
+*   [Install Postgres](https://cocoindex.io/docs/getting_started/installation#-install-postgres)
 *   Install [Neo4j](https://cocoindex.io/docs/targets/neo4j)
 *   [Configure your OpenAI API key](https://cocoindex.io/docs/ai/llm#openai).
 

{cocoindex-0.2.23 → cocoindex-0.3.1}/python/cocoindex/functions/colpali.py

@@ -2,7 +2,7 @@
 
 import functools
 from dataclasses import dataclass
-from typing import Any, Optional, TYPE_CHECKING, Literal
+from typing import Any, TYPE_CHECKING, Literal
 import numpy as np
 
 from .. import op
@@ -22,18 +22,11 @@ class ColPaliModelInfo:
     dimension: int
 
 
-@functools.lru_cache(maxsize=None)
+@functools.cache
 def _get_colpali_model_and_processor(model_name: str) -> ColPaliModelInfo:
     """Load and cache ColPali model and processor with shared device setup."""
     try:
-        from colpali_engine import (  # type: ignore[import-untyped]
-            ColPali,
-            ColPaliProcessor,
-            ColQwen2,
-            ColQwen2Processor,
-            ColSmol,
-            ColSmolProcessor,
-        )
+        import colpali_engine as ce  # type: ignore[import-untyped]
         import torch
     except ImportError as e:
         raise ImportError(
@@ -42,29 +35,30 @@ def _get_colpali_model_and_processor(model_name: str) -> ColPaliModelInfo:
         ) from e
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
+    lower_model_name = model_name.lower()
 
     # Determine model type from name
-    if "colpali" in model_name.lower():
-        model = ColPali.from_pretrained(
+    if lower_model_name.startswith("colpali"):
+        model = ce.ColPali.from_pretrained(
             model_name, torch_dtype=torch.bfloat16, device_map=device
         )
-        processor = ColPaliProcessor.from_pretrained(model_name)
-    elif "colqwen" in model_name.lower():
-        model = ColQwen2.from_pretrained(
+        processor = ce.ColPaliProcessor.from_pretrained(model_name)
+    elif lower_model_name.startswith("colqwen2.5"):
+        model = ce.ColQwen2_5.from_pretrained(
             model_name, torch_dtype=torch.bfloat16, device_map=device
         )
-        processor = ColQwen2Processor.from_pretrained(model_name)
-    elif "colsmol" in model_name.lower():
-        model = ColSmol.from_pretrained(
+        processor = ce.ColQwen2_5_Processor.from_pretrained(model_name)
+    elif lower_model_name.startswith("colqwen"):
+        model = ce.ColQwen2.from_pretrained(
             model_name, torch_dtype=torch.bfloat16, device_map=device
         )
-        processor = ColSmolProcessor.from_pretrained(model_name)
+        processor = ce.ColQwen2Processor.from_pretrained(model_name)
     else:
         # Fallback to ColPali for backwards compatibility
-        model = ColPali.from_pretrained(
+        model = ce.ColPali.from_pretrained(
             model_name, torch_dtype=torch.bfloat16, device_map=device
         )
-        processor = ColPaliProcessor.from_pretrained(model_name)
+        processor = ce.ColPaliProcessor.from_pretrained(model_name)
 
     # Detect dimension
     dimension = _detect_colpali_dimension(model, processor, device)
@@ -130,6 +124,7 @@ class ColPaliEmbedImage(op.FunctionSpec):
 @op.executor_class(
     gpu=True,
     cache=True,
+    batching=True,
     behavior_version=1,
 )
 class ColPaliEmbedImageExecutor:
@@ -146,7 +141,7 @@ class ColPaliEmbedImageExecutor:
         dimension = self._model_info.dimension
         return Vector[Vector[np.float32, Literal[dimension]]]  # type: ignore
 
-    def __call__(self, img_bytes: bytes) -> Any:
+    def __call__(self, img_bytes_list: list[bytes]) -> Any:
         try:
             from PIL import Image
             import torch
@@ -160,8 +155,11 @@ class ColPaliEmbedImageExecutor:
         processor = self._model_info.processor
         device = self._model_info.device
 
-        pil_image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
-        inputs = processor.process_images([pil_image]).to(device)
+        pil_images = [
+            Image.open(io.BytesIO(img_bytes)).convert("RGB")
+            for img_bytes in img_bytes_list
+        ]
+        inputs = processor.process_images(pil_images).to(device)
         with torch.no_grad():
             embeddings = model(**inputs)
 
@@ -171,10 +169,8 @@ class ColPaliEmbedImageExecutor:
                 f"Expected 3D tensor [batch, patches, hidden_dim], got shape {embeddings.shape}"
             )
 
-        # Keep patch-level embeddings: [batch, patches, hidden_dim] -> [patches, hidden_dim]
-        patch_embeddings = embeddings[0]  # Remove batch dimension
-
-        return patch_embeddings.cpu().to(torch.float32).numpy()
+        # [patches, hidden_dim]
+        return embeddings.cpu().to(torch.float32).numpy()
 
 
 class ColPaliEmbedQuery(op.FunctionSpec):
@@ -207,6 +203,7 @@ class ColPaliEmbedQuery(op.FunctionSpec):
     gpu=True,
     cache=True,
     behavior_version=1,
+    batching=True,
 )
 class ColPaliEmbedQueryExecutor:
     """Executor for ColVision query embedding (ColPali, ColQwen2, ColSmol, etc.)."""
@@ -222,7 +219,7 @@ class ColPaliEmbedQueryExecutor:
         dimension = self._model_info.dimension
         return Vector[Vector[np.float32, Literal[dimension]]]  # type: ignore
 
-    def __call__(self, query: str) -> Any:
+    def __call__(self, queries: list[str]) -> Any:
         try:
             import torch
         except ImportError as e:
@@ -234,7 +231,7 @@ class ColPaliEmbedQueryExecutor:
         processor = self._model_info.processor
         device = self._model_info.device
 
-        inputs = processor.process_queries([query]).to(device)
+        inputs = processor.process_queries(queries).to(device)
         with torch.no_grad():
             embeddings = model(**inputs)
 
@@ -244,7 +241,5 @@ class ColPaliEmbedQueryExecutor:
                 f"Expected 3D tensor [batch, tokens, hidden_dim], got shape {embeddings.shape}"
             )
 
-        # Keep token-level embeddings: [batch, tokens, hidden_dim] -> [tokens, hidden_dim]
-        token_embeddings = embeddings[0]  # Remove batch dimension
-
-        return token_embeddings.cpu().to(torch.float32).numpy()
+        # [tokens, hidden_dim]
+        return embeddings.cpu().to(torch.float32).numpy()

{cocoindex-0.2.23 → cocoindex-0.3.1}/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.1}/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}")