Code Operations

Relevant source files

The following files were used as context for generating this wiki page:

• docetl/operations/code_operations.py
• docetl/reasoning_optimizer/op_descriptions.py
• pyproject.toml
• tests/reasoning_optimizer/test_directive_apply.py
• docetl/reasoning_optimizer/directives/deterministic_doc_compression.py
• mkdocs.yml

Code Operations are a set of non-LLM-powered operators in DocETL that allow for deterministic data transformations using custom Python code. They are designed for tasks that do not require LLM reasoning, such as structured data manipulation, calculations, or regex-based processing. By executing pure Python, they offer a fast, cost-effective, and predictable way to process data within a pipeline.

There are three types of Code Operations: code_map, code_reduce, and code_filter. These operations are registered as plugins in the system via pyproject.toml, making them available to the DocETL runner. They are particularly useful in pipeline optimization, where expensive LLM-based operations can be swapped with more efficient code-based alternatives when the logic is deterministic.

Sources: docetl/reasoning_optimizer/op_descriptions.py:386-508, pyproject.toml:50-52

Core Concepts

Execution Model

All code operations leverage multi-threading for efficient processing. They use Python's concurrent.futures.ThreadPoolExecutor to apply the custom transform function to documents or groups of documents in parallel. The number of concurrent threads is configurable via the concurrent_thread_count parameter, which defaults to the number of CPU cores available on the machine.

Sources: docetl/operations/code_operations.py:31-34, docetl/operations/code_operations.py:108-111, docetl/operations/code_operations.py:178-181

Code Definition and Validation

The custom logic for each code operation is provided in the code parameter of its configuration. This parameter can be either a string containing the Python code or a callable Python function.

• If a string is provided, it must define a function named transform.
• If a callable is provided, the system uses inspect.getsource to retrieve its source code.

Before execution, a syntax_check is performed to ensure that the provided code defines a callable function named transform. This validation prevents runtime errors due to misconfigured code.

Sources: docetl/operations/code_operations.py:15-27, docetl/operations/code_operations.py:29-39

Class Structure

The three code operations share a common base structure, inheriting from BaseOperation. Each contains a schema inner class for Pydantic-based configuration validation and an execute method that implements the core logic.

The following diagram illustrates the class hierarchy for Code Operations.

Available Code Operations

CodeMapOperation

The code_map operation applies a Python function to each document in the input dataset independently. The function's return value, a dictionary, is merged with the original document.

Sources: docetl/reasoning_optimizer/op_descriptions.py:386-421, docetl/operations/code_operations.py:11-57

Execution Flow

This diagram shows the execution flow for a CodeMapOperation.

Configuration

Parameter	Type	Default	Description
name	string		A unique name for the operation.
type	string		Must be "code_map".
code	string or callable		Python code defining a transform(input_doc: dict) -> dict function. The returned dictionary is merged into the input document.
concurrent_thread_count	int	os.cpu_count()	The number of threads to use for concurrent execution.
drop_keys	list[str] or None	None	A list of keys to remove from the dictionary returned by the transform function before merging.

Sources: docetl/operations/code_operations.py:12-18, docetl/reasoning_optimizer/op_descriptions.py:392-398

CodeReduceOperation

The code_reduce operation groups documents by a specified reduce_key and applies a Python function to each group to produce a single aggregated output document.

Sources: docetl/operations/code_operations.py:60-138

Execution Flow

This diagram shows the execution flow for a CodeReduceOperation.

Configuration

Parameter	Type	Default	Description
name	string		A unique name for the operation.
type	string		Must be "code_reduce".
code	string or callable		Python code defining a transform(inputs: list[dict]) -> dict function that processes a group of documents and returns a single aggregated dictionary.
reduce_key	string or list[str]	"_all"	The key(s) to group documents by. If _all, all documents are processed as a single group.
concurrent_thread_count	int	os.cpu_count()	The number of threads to use for concurrent execution.

Sources: docetl/operations/code_operations.py:61-65

CodeFilterOperation

The code_filter operation filters documents based on a custom Python function. The function must return a boolean value; documents for which the function returns True are kept, while others are discarded.

Sources: docetl/reasoning_optimizer/op_descriptions.py:465-490, docetl/operations/code_operations.py:141-191

Execution Flow

This diagram shows the execution flow for a CodeFilterOperation.

Configuration

Parameter	Type	Default	Description
name	string		A unique name for the operation.
type	string		Must be "code_filter".
code	string or callable		Python code defining a transform(input_doc: dict) -> bool function. Returns True to keep the document, False to discard it.
concurrent_thread_count	int	os.cpu_count()	The number of threads to use for concurrent execution.

Sources: docetl/operations/code_operations.py:142-146, docetl/reasoning_optimizer/op_descriptions.py:471-477

Usage in Pipeline Optimization

Code Operations are integral to the docetl reasoning optimizer, which can automatically rewrite pipelines for better cost or accuracy. Directives within the optimizer often generate Code Operations to replace or augment more expensive LLM-based operations.

Deterministic Document Compression

The DeterministicDocCompressionDirective is used to reduce LLM processing costs by pre-processing documents. It inserts a code_map operation before an expensive downstream operation. This code_map uses deterministic logic, like regular expressions, to extract only the relevant parts of a document, thereby reducing the number of tokens sent to the LLM.

For example, it can generate a code_map to extract specific sections from a legal document before it is analyzed by an LLM.

Sources: docetl/reasoning_optimizer/directives/deterministic_doc_compression.py, tests/reasoning_optimizer/test_directive_apply.py:333-360

Swapping LLM Operations with Code

The SwapWithCodeDirective can replace an LLM-based reduce operation with a code_reduce operation. This is done when the aggregation logic is deterministic and can be expressed in Python, such as counting unique items. This swap significantly reduces cost and improves performance by avoiding an unnecessary LLM call.

The following test case demonstrates replacing an LLM-based reduce with a code_reduce to count unique locations.

# tests/reasoning_optimizer/test_directive_apply.py:619-646
def test_swap_with_code_apply():
    """Test that swap with code apply doesn't crash"""
    directive = SwapWithCodeDirective()
    
    # Test with reduce operation - no map needed (code reduce outputs correct schema)
    ops_list = [
        {
            "name": "count_locations",
            "type": "reduce",
            "reduce_key": "_all",
            "prompt": "Count unique locations: {% for input in inputs %}{{ input.location }}{% endfor %}",
            "model": "gpt-4o-mini",
            "output": {"schema": {"unique_count": "int", "locations": "list[str]"}}
        }
    ]
    
    from docetl.reasoning_optimizer.instantiate_schemas import SwapWithCodeInstantiateSchema
    rewrite_no_map = SwapWithCodeInstantiateSchema(
        code_reduce_name="code_count_locations",
        code='''def transform(inputs):
    locations = set()
    for item in inputs:
        if "location" in item:
            locations.add(item["location"])
    return {"unique_count": len(locations), "locations": sorted(list(locations))}''',
        map_prompt=None  # No map needed
    )
    
    result = directive.apply("gpt-4o-mini", ops_list, ["count_locations"], rewrite_no_map)
    assert isinstance(result, list)
    assert len(result) == 1  # Should replace reduce with code_reduce only
    assert result[0]["name"] == "count_locations"  # Should preserve original name
    assert result[0]["type"] == "code_reduce"

Sources: tests/reasoning_optimizer/test_directive_apply.py:619-646

Summary

Code Operations provide a powerful mechanism for integrating custom, deterministic Python logic into DocETL pipelines. They serve as a crucial tool for both manual pipeline authoring and automated optimization, enabling efficient data transformations, calculations, and filtering without the overhead and cost of LLM calls. Their multi-threaded execution and flexible code definition make them a versatile component for building robust and performant data processing workflows.