This is an automated email from the ASF dual-hosted git repository.
tballison pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tika.git
The following commit(s) were added to refs/heads/main by this push:
new 1f6e2782ae TIKA-4970: write image and page embeddings once per
document (#3342)
1f6e2782ae is described below
commit 1f6e2782ae2994d394e3d70a857fc11b1c6859fc
Author: Tim Allison <[email protected]>
AuthorDate: Mon Oct 19 15:35:22 2026 +0100
TIKA-4970: write image and page embeddings once per document (#3342)
---
CHANGES.txt | 4 ++
.../org/apache/tika/inference/ChunkTarget.java | 7 +-
.../java/org/apache/tika/inference/EmbedTask.java | 75 ++++++++++++++--------
.../org/apache/tika/inference/EmbedTaskTest.java | 66 +++++++++++++++++++
4 files changed, 125 insertions(+), 27 deletions(-)
diff --git a/CHANGES.txt b/CHANGES.txt
index 4b04bc03bc..4946d32dbf 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -162,6 +162,10 @@ Release 4.2.0 - unreleased
MetadataFilter; 4.1.0 failed to install on its own in Equinox/p2
(TIKA-4945).
+ * The embed task writes a document's image and page embeddings to
+ tk:chunks once per binding instead of once per image or page, which
+ re-serialized the whole array each time (TIKA-4970).
+
Release 4.1.0 - 9/26/2026
HIGHLIGHTS
diff --git
a/tika-parsers/tika-parsers-ml/tika-inference/src/main/java/org/apache/tika/inference/ChunkTarget.java
b/tika-parsers/tika-parsers-ml/tika-inference/src/main/java/org/apache/tika/inference/ChunkTarget.java
index a279552396..de869b1939 100644
---
a/tika-parsers/tika-parsers-ml/tika-inference/src/main/java/org/apache/tika/inference/ChunkTarget.java
+++
b/tika-parsers/tika-parsers-ml/tika-inference/src/main/java/org/apache/tika/inference/ChunkTarget.java
@@ -82,6 +82,11 @@ public final class ChunkTarget {
/** Tags the chunks with the locator, if any, and appends them to the
field. */
public void write(List<Chunk> chunks, String field) throws IOException {
+ ChunkSerializer.mergeInto(metadata, tag(chunks), field);
+ }
+
+ /** Tags the chunks with the locator, if any, for a caller that writes
them itself. */
+ List<Chunk> tag(List<Chunk> chunks) {
if (locator != null) {
for (Chunk chunk : chunks) {
if (chunk.getLocators() != null) {
@@ -89,6 +94,6 @@ public final class ChunkTarget {
}
}
}
- ChunkSerializer.mergeInto(metadata, chunks, field);
+ return chunks;
}
}
diff --git
a/tika-parsers/tika-parsers-ml/tika-inference/src/main/java/org/apache/tika/inference/EmbedTask.java
b/tika-parsers/tika-parsers-ml/tika-inference/src/main/java/org/apache/tika/inference/EmbedTask.java
index e23f9e2448..b9aba497ee 100644
---
a/tika-parsers/tika-parsers-ml/tika-inference/src/main/java/org/apache/tika/inference/EmbedTask.java
+++
b/tika-parsers/tika-parsers-ml/tika-inference/src/main/java/org/apache/tika/inference/EmbedTask.java
@@ -99,31 +99,58 @@ public class EmbedTask implements InferenceTask {
}
int batchSize = Math.max(1, embedder.getMaxBatchSize());
TikaException first = null;
- for (int start = 0; start < units.size(); start += batchSize) {
- List<InferenceUnit> batch = units.subList(start, Math.min(start +
batchSize, units.size()));
- try {
- embed(binding, batch, embedder, context);
- } catch (TikaException e) {
- if (batch.size() == 1) {
- first = first == null ? e : first;
- continue;
- }
- // one bad image fails a whole batch; retry singly so the rest
keep their vectors
- for (InferenceUnit unit : batch) {
- try {
- embed(binding, List.of(unit), embedder, context);
- } catch (TikaException single) {
- first = first == null ? single : first;
+ // written once per destination at the end: each write re-serializes
the whole field
+ List<Owned> embedded = new ArrayList<>();
+ try {
+ for (int start = 0; start < units.size(); start += batchSize) {
+ List<InferenceUnit> batch = units.subList(start,
Math.min(start + batchSize, units.size()));
+ try {
+ embed(binding, batch, embedder, context, embedded);
+ } catch (TikaException e) {
+ if (batch.size() == 1) {
+ first = first == null ? e : first;
+ continue;
+ }
+ // one bad image fails a whole batch; retry singly so the
rest keep their vectors
+ for (InferenceUnit unit : batch) {
+ try {
+ embed(binding, List.of(unit), embedder, context,
embedded);
+ } catch (TikaException single) {
+ first = first == null ? single : first;
+ }
}
}
}
+ } catch (IOException | RuntimeException e) {
+ write(embedded, e);
+ throw e;
}
+ write(embedded, null);
if (first != null) {
throw first;
}
}
- /** A chunk of text and the document it belongs to. */
+ /** Writes each destination's chunks at once; a write failure rides on the
failure already in flight. */
+ private static void write(List<Owned> embedded, Exception inFlight) throws
IOException {
+ for (List<Owned> byOwner : byOwner(embedded)) {
+ List<Chunk> chunks = new ArrayList<>(byOwner.size());
+ for (Owned o : byOwner) {
+ chunks.add(o.chunk());
+ }
+ try {
+ ChunkTarget.self(byOwner.get(0).owner())
+ .write(chunks,
TikaCoreProperties.TIKA_CHUNKS.getName());
+ } catch (IOException e) {
+ if (inFlight == null) {
+ throw e;
+ }
+ inFlight.addSuppressed(e);
+ }
+ }
+ }
+
+ /** A chunk and the document it is written to. */
private record Owned(Chunk chunk, Metadata owner) {
}
@@ -174,14 +201,7 @@ public class EmbedTask implements InferenceTask {
}
}
}
- for (List<Owned> byOwner : byOwner(embedded)) {
- List<Chunk> chunks = new ArrayList<>(byOwner.size());
- for (Owned o : byOwner) {
- chunks.add(o.chunk());
- }
- ChunkTarget.self(byOwner.get(0).owner())
- .write(chunks, TikaCoreProperties.TIKA_CHUNKS.getName());
- }
+ write(embedded, null);
if (first != null) {
throw first;
}
@@ -382,8 +402,9 @@ public class EmbedTask implements InferenceTask {
return chunks;
}
+ /** Embeds the batch and adds each unit's chunk, tagged for its
destination, to {@code embedded}. */
private static void embed(InferenceBinding binding, List<InferenceUnit>
batch,
- EmbeddingEngine embedder, ParseContext context)
+ EmbeddingEngine embedder, ParseContext context,
List<Owned> embedded)
throws IOException, TikaException {
List<byte[]> images = new ArrayList<>(batch.size());
List<String> mimeTypes = new ArrayList<>(batch.size());
@@ -405,7 +426,9 @@ public class EmbedTask implements InferenceTask {
chunk.setVector(vectors.get(i));
chunk.setProducer(binding.getId());
chunk.setModality(binding.getModality());
- ChunkTarget.of(unit).write(List.of(chunk),
TikaCoreProperties.TIKA_CHUNKS.getName());
+ ChunkTarget target = ChunkTarget.of(unit);
+ target.tag(List.of(chunk));
+ embedded.add(new Owned(chunk, target.getMetadata()));
}
}
}
diff --git
a/tika-parsers/tika-parsers-ml/tika-inference/src/test/java/org/apache/tika/inference/EmbedTaskTest.java
b/tika-parsers/tika-parsers-ml/tika-inference/src/test/java/org/apache/tika/inference/EmbedTaskTest.java
index 72143330d2..c6d824c7e6 100644
---
a/tika-parsers/tika-parsers-ml/tika-inference/src/test/java/org/apache/tika/inference/EmbedTaskTest.java
+++
b/tika-parsers/tika-parsers-ml/tika-inference/src/test/java/org/apache/tika/inference/EmbedTaskTest.java
@@ -24,7 +24,9 @@ import static org.junit.jupiter.api.Assertions.assertTrue;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
+import java.util.Arrays;
import java.util.List;
+import java.util.Map;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
@@ -38,6 +40,7 @@ import org.apache.tika.http.TikaTestHttpServer;
import org.apache.tika.metadata.Metadata;
import org.apache.tika.metadata.TikaCoreProperties;
import org.apache.tika.metadata.TikaPagedText;
+import org.apache.tika.metadata.writelimiter.MetadataWriteLimiter;
import org.apache.tika.mime.MediaType;
import org.apache.tika.parser.ParseContext;
import org.apache.tika.parser.inference.Engine;
@@ -126,6 +129,56 @@ public class EmbedTaskTest {
}
}
+ /** Each write re-serializes the field, so a destination is written once,
not once per unit. */
+ @Test
+ public void testOneWritePerDestination() throws Exception {
+ server.enqueue(new TikaTestHttpServer.MockResponse(200, response(0,
1)));
+ server.enqueue(new TikaTestHttpServer.MockResponse(200, response(0,
1)));
+ server.enqueue(new TikaTestHttpServer.MockResponse(200, response(0)));
+ int[] parentWrites = new int[1];
+ int[] otherWrites = new int[1];
+ Metadata parent = new Metadata(countingChunkWrites(parentWrites));
+ Metadata other = new Metadata(countingChunkWrites(otherWrites));
+ List<InferenceUnit> units = List.of(unit(parent, "a", null),
unit(other, "b", null),
+ unit(parent, "c", null), unit(parent, "d", null), unit(other,
"e", null));
+
+ new EmbedTask().run(binding(), units, engine, new ParseContext());
+
+ assertEquals(3, server.getRequestCount());
+ assertEquals(1, parentWrites[0]);
+ assertEquals(1, otherWrites[0]);
+ List<Chunk> chunks =
ChunkSerializer.fromJson(parent.get(TikaCoreProperties.TIKA_CHUNKS));
+ assertEquals(3, chunks.size());
+ assertEquals("/a",
chunks.get(0).getLocators().getEmbedded().get(0).getIdPath());
+ assertEquals("/c",
chunks.get(1).getLocators().getEmbedded().get(0).getIdPath());
+ assertEquals("/d",
chunks.get(2).getLocators().getEmbedded().get(0).getIdPath());
+ assertEquals(2,
ChunkSerializer.fromJson(other.get(TikaCoreProperties.TIKA_CHUNKS)).size());
+ }
+
+ private static MetadataWriteLimiter countingChunkWrites(int[] count) {
+ return new MetadataWriteLimiter() {
+ @Override
+ public void add(String field, String value, Map<String, String[]>
data) {
+ String[] values = data.get(field);
+ if (values == null) {
+ set(field, value, data);
+ return;
+ }
+ String[] appended = Arrays.copyOf(values, values.length + 1);
+ appended[values.length] = value;
+ data.put(field, appended);
+ }
+
+ @Override
+ public void set(String field, String value, Map<String, String[]>
data) {
+ if (TikaCoreProperties.TIKA_CHUNKS.getName().equals(field)) {
+ count[0]++;
+ }
+ data.put(field, new String[]{value});
+ }
+ };
+ }
+
@Test
public void testNoParentKeepsTheChunk() throws Exception {
server.enqueue(new TikaTestHttpServer.MockResponse(200, response(0)));
@@ -251,6 +304,19 @@ public class EmbedTaskTest {
"the good unit keeps its vector");
}
+ @Test
+ public void testFailureKeepsWhatWasEmbedded() throws Exception {
+ server.enqueue(new TikaTestHttpServer.MockResponse(200, response(0,
1)));
+ server.enqueue(new TikaTestHttpServer.MockResponse(200, response(0)));
+ Metadata parent = new Metadata();
+ List<InferenceUnit> units = List.of(unit(parent, "a", null),
unit(parent, "b", null),
+ unit(parent, "c", "not-a-page"));
+ assertThrows(NumberFormatException.class,
+ () -> new EmbedTask().run(binding(), units, engine, new
ParseContext()));
+ assertEquals(2,
ChunkSerializer.fromJson(parent.get(TikaCoreProperties.TIKA_CHUNKS)).size(),
+ "the first batch is written before the failure propagates");
+ }
+
@Test
public void testResponseCountMustMatch() throws Exception {
server.enqueue(new TikaTestHttpServer.MockResponse(200, response(0)));
