Advanced Operations: Cluster, Rank, EquiJoin

Relevant source files

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

• docetl/operations/cluster.py
• docetl/operations/rank.py
• docetl/operations/equijoin.py
• tests/ranking/test_rank.py
• pyproject.toml
• tests/ranking/test_rank_budget.py

The docetl framework provides several advanced operations for complex data manipulation tasks beyond simple mapping and filtering. This page details three such operations: cluster, rank, and equijoin. These operations are registered as official plugins within the docetl ecosystem.

The rank operation enables sophisticated, LLM-driven ordering of documents based on nuanced criteria. The equijoin operation provides a classic database-style join to merge two datasets on a common key. The cluster operation is defined but currently acts as a pass-through, returning the input data without modification.

Sources: pyproject.toml:41-43

Rank Operation

The rank operation sorts a list of documents based on a user-defined prompt. It employs a variety of internal strategies, including an advanced "Picky" method that uses an initial ordering followed by iterative LLM-based pairwise comparisons to refine the ranking. This operation is designed for tasks where simple sorting is insufficient and nuanced, contextual understanding is required.

Sources: docetl/operations/rank.py, tests/ranking/test_rank.py

Configuration

The rank operation is configured with the following parameters in the pipeline YAML:

Parameter	Type	Description	Default
name	string	A unique name for the operation.	N/A
type	string	Must be order.	order
prompt	string	The LLM prompt describing the ranking criteria.	N/A
input_keys	list[string]	The document keys to use as input for the prompt.	N/A
direction	string	The sort direction, either asc (ascending) or desc (descending).	asc
batch_size	int	The number of documents to process in a single LLM call for certain ranking methods.	5
initial_ordering_method	string	The method for the initial sort. Can be embedding, likert, or calibrated_embedding.	embedding
rerank_call_budget	int	The maximum number of LLM calls for the re-ranking phase.	20
k	int	The number of top/bottom items to focus on during ranking.	10
verbose	bool	If True, enables detailed logging.	False

Sources: docetl/operations/rank.py:27-46, tests/ranking/test_rank.py:60-75

Execution Flow (execute method)

The primary execution logic, also referred to as the "Picky" method in tests, follows a multi-step process to achieve a refined ranking. It combines an efficient initial sort with a budget-constrained, high-quality LLM re-ranking phase.

This diagram illustrates the flow of the main execute method in the RankOperation class.

1. Initial Ordering: The operation begins by creating a preliminary sorted list using one of three methods specified by initial_ordering_method.
   • embedding (default): Sorts documents based on embedding similarity.
   • likert: Sorts documents by assigning them a Likert scale value via an LLM call.
   • calibrated_embedding: A variation of embedding sort.
2. Create Reranking Pairs: It generates pairs of documents for pairwise comparison, focusing on items where the initial ranking may be ambiguous.
3. Iterative Reranking: Within a rerank_call_budget, it sends pairs to an LLM to determine which of the two has a higher rank according to the prompt.
4. Graph Construction: The pairwise preferences are used to build a directed graph where an edge A -> B means A is ranked higher than B.
5. Topological Sort: A topological sort of the graph produces the final, refined ranking.
6. Output: The documents are returned with a _rank field added.

Sources: docetl/operations/rank.py:53-156

Internal Ranking Strategies

The RankOperation class contains several internal methods for different ranking approaches, which are used for testing and comparison.

Method	Description
_execute_comparison_qurk	A baseline method that performs all-pairs comparison. It can be very expensive.
_execute_rating_embedding_qurk	Ranks documents by generating an embedding for each and sorting by similarity to the prompt.
_execute_likert_rating_qurk	Ranks documents by asking an LLM to assign a Likert scale rating to each one.
_execute_sliding_window_qurk	Performs pairwise comparisons in a sliding window over an initially sorted list.
_execute_calibrated_embedding_sort	An alternative embedding-based sorting method.

Sources: tests/ranking/test_rank.py:130-170, docetl/operations/rank.py:158-403

Performance Evaluation

The repository's tests evaluate ranking methods using standard information retrieval metrics:

• Kendall's Tau: Measures the ordinal association between the method's ranking and a ground-truth ranking.
• Normalized Discounted Cumulative Gain (NDCG): Evaluates the quality of the ranking at the top k results, considering the relevance of each item.

Tests are performed on various datasets, including synthetic squares, medical transcripts, and scientific abstracts, to compare the cost, runtime, and accuracy of different ranking strategies.

Sources: tests/ranking/test_rank.py:23-31, 202-212, tests/ranking/test_rank_budget.py:22-43

EquiJoin Operation

The equijoin operation performs an inner join on two datasets using specified keys. It merges documents from a "left" and "right" dataset where the values of the left_on and right_on keys match.

Sources: docetl/operations/equijoin.py

Configuration

Parameter	Type	Description	Default
name	string	A unique name for the operation.	N/A
type	string	Must be equijoin.	equijoin
right_input	string	The name of the "right" dataset to join with.	N/A
left_on	string	The key from the "left" (primary) dataset to join on.	N/A
right_on	string	The key from the "right" dataset to join on.	N/A

Sources: docetl/operations/equijoin.py:15-22

Execution Flow

The join is executed by first processing the "right" dataset into a lookup structure and then iterating through the "left" dataset to find matches.

This diagram shows the process of the EquiJoinOperation.

1. Group Right Dataset: The operation first iterates through the right_input dataset and groups all documents by the value of the right_on key. This creates a dictionary where each key is a join value and the value is a list of all documents from the right dataset with that join value.
2. Iterate Left Dataset: It then iterates through each document in the "left" (primary) dataset.
3. Find and Merge: For each left document, it looks up the value of its left_on key in the grouped right dataset dictionary.
4. Create Results: If a match is found, it creates new documents by merging the left document with each of the corresponding right documents. If no match is found, the left document is dropped (behavior of an inner join).
5. Return: The final output is a new dataset containing only the merged documents.

Sources: docetl/operations/equijoin.py:24-49

Cluster Operation

The cluster operation is defined as an available operation within the docetl framework. However, its current implementation is a placeholder.

Sources: docetl/operations/cluster.py, pyproject.toml:44

Configuration

Parameter	Type	Description	Default
name	string	A unique name for the operation.	N/A
type	string	Must be cluster.	cluster

Sources: docetl/operations/cluster.py:12-16

Execution Flow

The execute method for the ClusterOperation currently returns the input data without any modification. It acts as a simple pass-through.

Conclusion

The docetl library offers powerful, high-level operations for sophisticated data processing pipelines. The rank operation provides a flexible and robust solution for LLM-based sorting with multiple strategies to balance cost and accuracy. The equijoin operation enables relational data merging, essential for integrating information from different sources. While the cluster operation is defined, its functionality is not yet implemented, suggesting it may be a feature planned for future development. These advanced operations significantly extend the capabilities of docetl beyond basic document transformation.