Skip to content

TokenCountBatchingStrategy silently drops equal documents from batches #6560

Description

@subhashpolisetti

Bug description

TokenCountBatchingStrategy.batch(List<Document>) collects per-document token
counts into a Map<Document, Integer> before forming batches. Since
Document.equals() compares id, text, media, metadata and score, two equal
documents collapse into a single map entry, so the duplicate is silently
dropped from the returned batches.

batch is a partitioning operation, and its only caller,
EmbeddingModel.embed(List<Document>, ...), maps the returned embeddings back
to the input documents by position and asserts
embeddings.size() == documents.size(). When a duplicate is dropped, the
provider returns fewer embeddings than input documents and ingestion fails
with:

IllegalArgumentException: Embeddings must have the same number as that of the documents

which points at the embedding provider rather than the batching step.

This happens in practice with content-based id generators (e.g.
JdkSha256HexIdGenerator), where identical chunks such as repeated
boilerplate produce equal Documents, and trivially when the same Document
instance appears twice in the list.

Environment

Reproduced on Spring AI main. Plain JVM, no provider involved
the failing path is entirely within TokenCountBatchingStrategy, so it
reproduces offline.

Steps to reproduce

TokenCountBatchingStrategy strategy = new TokenCountBatchingStrategy();
Document d1 = new Document("shared-id", "same text content", Map.of());
Document d2 = new Document("shared-id", "same text content", Map.of());

List<List<Document>> batches = strategy.batch(List.of(d1, d2));
// batches contain 1 document in total, not 2

The same happens with strategy.batch(List.of(d1, d1)).

Expected behavior

batch preserves every input document, regardless of document equality.
Deduplication policy, if desired, belongs to the ingestion layer that owns
document identity, not to the token-count partitioner.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions