From 9f816c3503b82a18df6dbeac0efc82fdf6ddce76 Mon Sep 17 00:00:00 2001 From: Koustubh Gupta Date: Thu, 18 Jun 2026 09:40:52 +0530 Subject: [PATCH 1/2] Add get-by-id via DocumentLookupProvider SPI in DataFormatAwareEngine. Add version map and integrated get-by-id for document lookup. Signed-off-by: Koustubh Gupta --- .../analytics/spi/DocumentLookupService.java | 150 ++++++ .../analytics/spi/DocumentRowReader.java | 51 ++ .../analytics/spi/ReservedFields.java | 36 ++ .../spi/DocumentLookupServiceTests.java | 267 ++++++++++ .../be/datafusion/DataFusionPlugin.java | 62 ++- .../opensearch/be/datafusion/GetService.java | 245 +++++++++ .../be/datafusion/GetServiceTests.java | 155 ++++++ .../be/lucene/LuceneDocumentResolver.java | 106 ++++ .../opensearch/be/lucene/LucenePlugin.java | 19 +- .../lucene/LuceneDocumentResolverTests.java | 175 ++++++ .../analytics/exec/ArrowValues.java | 72 +++ .../analytics/exec/ArrowValuesTests.java | 116 ++++ sandbox/plugins/composite-engine/build.gradle | 1 + .../DataFormatAwareReadonlyEngineBaseIT.java | 4 +- .../index/engine/DataFormatAwareEngine.java | 254 ++++++++- .../DataFormatAwareNRTReplicationEngine.java | 27 + .../engine/DataFormatAwareReadOnlyEngine.java | 25 + .../index/engine/EngineBackedIndexer.java | 8 + .../opensearch/index/engine/EngineConfig.java | 38 +- .../index/engine/EngineConfigFactory.java | 35 +- .../engine/exec/DocumentLookupSupport.java | 103 ++++ .../engine/exec/DocumentMetadataResolver.java | 55 ++ .../opensearch/index/engine/exec/Indexer.java | 9 + .../engine/exec/coord/CatalogSnapshot.java | 18 + .../index/get/DocumentLookupResult.java | 141 +++++ .../opensearch/index/get/ShardGetService.java | 92 ++-- .../opensearch/index/shard/IndexShard.java | 6 +- .../opensearch/indices/IndicesService.java | 26 +- .../plugins/DocumentLookupProvider.java | 70 +++ .../engine/DataFormatAwareEngineTests.java | 504 +++++++++++++++++- .../DataFormatAwareReadOnlyEngineTests.java | 137 +++++ .../DataformatAwareCatalogSnapshotTests.java | 41 ++ .../index/get/DocumentLookupResultTests.java | 68 +++ .../index/shard/ShardGetServiceTests.java | 60 +++ .../plugins/DocumentLookupProviderTests.java | 34 ++ 35 files changed, 3162 insertions(+), 48 deletions(-) create mode 100644 sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/DocumentLookupService.java create mode 100644 sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/DocumentRowReader.java create mode 100644 sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ReservedFields.java create mode 100644 sandbox/libs/analytics-framework/src/test/java/org/opensearch/analytics/spi/DocumentLookupServiceTests.java create mode 100644 sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/GetService.java create mode 100644 sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/GetServiceTests.java create mode 100644 sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneDocumentResolver.java create mode 100644 sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneDocumentResolverTests.java create mode 100644 server/src/main/java/org/opensearch/index/engine/exec/DocumentLookupSupport.java create mode 100644 server/src/main/java/org/opensearch/index/engine/exec/DocumentMetadataResolver.java create mode 100644 server/src/main/java/org/opensearch/index/get/DocumentLookupResult.java create mode 100644 server/src/main/java/org/opensearch/plugins/DocumentLookupProvider.java create mode 100644 server/src/test/java/org/opensearch/index/get/DocumentLookupResultTests.java create mode 100644 server/src/test/java/org/opensearch/plugins/DocumentLookupProviderTests.java diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/DocumentLookupService.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/DocumentLookupService.java new file mode 100644 index 0000000000000..ea6dd83d4835e --- /dev/null +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/DocumentLookupService.java @@ -0,0 +1,150 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.analytics.spi; + +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.common.lucene.uid.Versions; +import org.opensearch.common.xcontent.XContentFactory; +import org.opensearch.core.common.bytes.BytesReference; +import org.opensearch.core.index.Index; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.index.engine.exec.DocumentMetadataResolver; +import org.opensearch.index.engine.exec.IndexReaderProvider; +import org.opensearch.index.engine.exec.Segment; +import org.opensearch.index.engine.exec.WriterFileSet; +import org.opensearch.index.get.DocumentLookupResult; +import org.opensearch.index.seqno.SequenceNumbers; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +/** + * Core orchestrator for document lookup. Coordinates row-location resolution, + * backend-specific execution, and result assembly. + * + * @opensearch.experimental + */ +@ExperimentalApi +public class DocumentLookupService { + + private final DocumentMetadataResolver documentResolver; + private final DocumentRowReader executor; + + public DocumentLookupService(DocumentMetadataResolver documentResolver, DocumentRowReader executor) { + this.documentResolver = documentResolver; + this.executor = executor; + } + + public DocumentLookupResult getById(String id, IndexReaderProvider.Reader reader, Index index) throws IOException { + Map row = fetchRow(id, reader); + return row == null ? DocumentLookupResult.notFound(id) : buildResultFromRow(id, row); + } + + /** + * Resolves only the version metadata ({@code _version}/{@code _seq_no}/{@code _primary_term}) for an id, + * skipping {@code _source} reconstruction. + */ + public DocumentLookupResult getVersionMetadata(String id, IndexReaderProvider.Reader reader, Index index) throws IOException { + Map row = fetchRow(id, reader); + if (row == null) { + return DocumentLookupResult.notFound(id); + } + long seqNo = extractLong(row, "_seq_no", SequenceNumbers.UNASSIGNED_SEQ_NO); + long primaryTerm = extractLong(row, "_primary_term", SequenceNumbers.UNASSIGNED_PRIMARY_TERM); + long version = extractLong(row, "_version", Versions.NOT_FOUND); + return new DocumentLookupResult(id, version, true, null, seqNo, primaryTerm, Map.of(), Map.of()); + } + + /** Locates an id via the resolver and fetches its raw row. Returns null only when the id is not found. */ + private Map fetchRow(String id, IndexReaderProvider.Reader reader) throws IOException { + DocumentMetadataResolver.DocumentMetadata metadata = documentResolver.resolveMetadata(reader, id); + if (metadata == null) { + return null; + } + + WriterFileSet fileSet = reader.catalogSnapshot().findFileSet(executor.formatName(), metadata.writerGeneration()); + if (fileSet == null) { + throw new IllegalStateException( + "Resolver located id [" + + id + + "] at writer generation [" + + metadata.writerGeneration() + + "] but no matching file set was found" + ); + } + + Map row = executor.executeSingleRow(metadata.rowId(), fileSet); + if (row == null) { + throw new IllegalStateException( + "Resolver located id [" + + id + + "] at writer generation [" + + metadata.writerGeneration() + + "] rowId [" + + metadata.rowId() + + "] but backend returned no row" + ); + } + return row; + } + + public List getDocsAboveSeqNo(long fromSeqNoExclusive, IndexReaderProvider.Reader reader, Index index) + throws IOException { + List fileSets = new ArrayList<>(); + for (Segment segment : reader.catalogSnapshot().getSegments()) { + WriterFileSet fileSet = segment.dfGroupedSearchableFiles().get(executor.formatName()); + if (fileSet != null && !fileSet.files().isEmpty()) { + fileSets.add(fileSet); + } + } + List results = new ArrayList<>(); + for (Map row : executor.executeRowsAboveSeqNo(fileSets, fromSeqNoExclusive)) { + Object idVal = row.get("_id"); + if (idVal != null) { + results.add(buildResultFromRow(idVal.toString(), row)); + } + } + return results; + } + + /** Builds a DocumentLookupResult from a raw row, filtering reserved/internal fields. */ + private static DocumentLookupResult buildResultFromRow(String id, Map row) throws IOException { + long seqNo = extractLong(row, "_seq_no", SequenceNumbers.UNASSIGNED_SEQ_NO); + long primaryTerm = extractLong(row, "_primary_term", SequenceNumbers.UNASSIGNED_PRIMARY_TERM); + long version = extractLong(row, "_version", Versions.NOT_FOUND); + + Map filtered = new LinkedHashMap<>(); + for (Map.Entry e : row.entrySet()) { + if (ReservedFields.RESERVED.contains(e.getKey()) || e.getKey().startsWith("_")) continue; + filtered.put(e.getKey(), e.getValue()); + } + + BytesReference source; + try (XContentBuilder xcb = XContentFactory.jsonBuilder()) { + xcb.map(filtered); + source = BytesReference.bytes(xcb); + } + + return new DocumentLookupResult(id, version, true, source, seqNo, primaryTerm, Map.of(), Map.of()); + } + + public static long extractLong(Map row, String key, long fallback) { + Object v = row.get(key); + if (v == null) return fallback; + if (v instanceof Number) return ((Number) v).longValue(); + try { + return Long.parseLong(v.toString()); + } catch (NumberFormatException e) { + return fallback; + } + } +} diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/DocumentRowReader.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/DocumentRowReader.java new file mode 100644 index 0000000000000..7297dc9439bb1 --- /dev/null +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/DocumentRowReader.java @@ -0,0 +1,51 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.analytics.spi; + +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.index.engine.exec.WriterFileSet; + +import java.io.IOException; +import java.util.List; +import java.util.Map; + +/** + * Backend-specific execution contract for reading document rows. Backends implement this + * to perform the actual storage read (e.g., DataFusion native parquet scan) from a file set + * that the Core layer has already resolved from the catalog snapshot. + * + * @opensearch.experimental + */ +@ExperimentalApi +public interface DocumentRowReader { + + /** + * The storage format this backend reads (e.g. {@code "parquet"}). Lets the Core layer + * resolve the backend's candidate {@link WriterFileSet} from the catalog snapshot. + */ + String formatName(); + + /** + * Fetch a single row at the given offset from the pre-resolved file set. + * + * @param rowId the row offset to fetch + * @param fileSet the file set to read from + * @return the row as a field-name → value map, or null if not found + */ + Map executeSingleRow(long rowId, WriterFileSet fileSet) throws IOException; + + /** + * Fetch all rows with {@code _seq_no > fromSeqNoExclusive} from the Core-resolved file sets + * (one per segment for this backend's format). + * + * @param fileSets the file sets to scan + * @param fromSeqNoExclusive the exclusive lower bound on {@code _seq_no} + */ + List> executeRowsAboveSeqNo(List fileSets, long fromSeqNoExclusive) throws IOException; +} diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ReservedFields.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ReservedFields.java new file mode 100644 index 0000000000000..a3d8746ca0ffa --- /dev/null +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ReservedFields.java @@ -0,0 +1,36 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.analytics.spi; + +import java.util.Set; + +/** + * Metadata fields managed by the storage layer that must not appear in + * user-visible {@code _source}. Backends use this set to filter internal + * columns from document lookup results. + * + * @opensearch.internal + */ +public final class ReservedFields { + + private ReservedFields() {} + + /** Fields reserved by the engine/storage layer — excluded from _source reconstruction. */ + public static final Set RESERVED = Set.of( + "_id", + "_seq_no", + "_primary_term", + "_version", + "_doc_count", + "_size", + "_routing", + "_ignored", + "__row_id__" + ); +} diff --git a/sandbox/libs/analytics-framework/src/test/java/org/opensearch/analytics/spi/DocumentLookupServiceTests.java b/sandbox/libs/analytics-framework/src/test/java/org/opensearch/analytics/spi/DocumentLookupServiceTests.java new file mode 100644 index 0000000000000..22762991c22d2 --- /dev/null +++ b/sandbox/libs/analytics-framework/src/test/java/org/opensearch/analytics/spi/DocumentLookupServiceTests.java @@ -0,0 +1,267 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.analytics.spi; + +import org.opensearch.common.lucene.uid.Versions; +import org.opensearch.core.index.Index; +import org.opensearch.index.engine.exec.DocumentMetadataResolver; +import org.opensearch.index.engine.exec.IndexReaderProvider; +import org.opensearch.index.engine.exec.Segment; +import org.opensearch.index.engine.exec.WriterFileSet; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; +import org.opensearch.index.get.DocumentLookupResult; +import org.opensearch.index.seqno.SequenceNumbers; +import org.opensearch.test.OpenSearchTestCase; + +import java.nio.file.Path; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +import org.mockito.ArgumentCaptor; + +import static org.mockito.ArgumentMatchers.anyList; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** + * Unit tests for {@link DocumentLookupService}, the Core orchestrator for document lookup. + * + *

All collaborators are interfaces ({@link DocumentMetadataResolver}, {@link DocumentRowReader}, + * {@link IndexReaderProvider.Reader}) or an abstract class ({@link CatalogSnapshot}) and are mocked. + * {@link Segment}/{@link WriterFileSet} are real objects constructed via their builders since they + * are final records / sealed classes. + * + *

Coverage: + *

    + *
  • {@code getById} — not-found, success (reserved/underscore field filtering), and the two + * Lucene/Parquet inconsistency ISE paths (missing file set, missing row).
    • + *
  • {@code getVersionMetadata} — not-found, version-field extraction with null source, and + * fallback defaults when version fields are absent.
    • + *
  • {@code getDocsAboveSeqNo} — rows without {@code _id} skipped, empty/wrong-format file sets + * excluded from the backend scan, and empty result handling.
    • + *
  • {@code extractLong} — Number, missing key, parsable String, and unparsable String.
    • + *
+ */ +public class DocumentLookupServiceTests extends OpenSearchTestCase { + + private static final String FORMAT = "parquet"; + + private static final Index INDEX = new Index("idx", "uuid"); + + private DocumentMetadataResolver resolver; + private DocumentRowReader executor; + private IndexReaderProvider.Reader reader; + private CatalogSnapshot snapshot; + private DocumentLookupService service; + + @Override + public void setUp() throws Exception { + super.setUp(); + resolver = mock(DocumentMetadataResolver.class); + executor = mock(DocumentRowReader.class); + reader = mock(IndexReaderProvider.Reader.class); + snapshot = mock(CatalogSnapshot.class); + when(reader.catalogSnapshot()).thenReturn(snapshot); + when(executor.formatName()).thenReturn(FORMAT); + service = new DocumentLookupService(resolver, executor); + } + + // ---- helpers ------------------------------------------------------------ + + private static WriterFileSet fileSet(long generation, String... files) { + WriterFileSet.Builder b = WriterFileSet.builder() + .directory(Path.of("/tmp/idx")) + .writerGeneration(generation) + .addNumRows(files.length); + for (String f : files) { + b.addFile(f); + } + return b.build(); + } + + private static DocumentMetadataResolver.DocumentMetadata metadata(String id, long rowId, long generation) { + return new DocumentMetadataResolver.DocumentMetadata(id, rowId, generation); + } + + private static Map row(Object... kv) { + Map row = new LinkedHashMap<>(); + for (int i = 0; i < kv.length; i += 2) { + row.put((String) kv[i], kv[i + 1]); + } + return row; + } + + // ---- getById ------------------------------------------------------------ + + public void testGetById_notFoundWhenResolverReturnsNull() throws Exception { + when(resolver.resolveMetadata(reader, "missing")).thenReturn(null); + + DocumentLookupResult result = service.getById("missing", reader, INDEX); + + assertFalse(result.exists()); + assertEquals("missing", result.id()); + assertEquals(Versions.NOT_FOUND, result.version()); + assertNull(result.source()); + } + + public void testGetById_successFiltersReservedAndUnderscoreFields() throws Exception { + when(resolver.resolveMetadata(reader, "doc1")).thenReturn(metadata("doc1", 0L, 7L)); + WriterFileSet fs = fileSet(7L, "0.parquet"); + when(snapshot.findFileSet(FORMAT, 7L)).thenReturn(fs); + when(executor.executeSingleRow(0L, fs)).thenReturn( + row("_id", "doc1", "_seq_no", 42L, "_primary_term", 2L, "_version", 5L, "_custom", "internal", "name", "alice", "age", 30) + ); + + DocumentLookupResult result = service.getById("doc1", reader, INDEX); + + assertTrue(result.exists()); + assertEquals("doc1", result.id()); + assertEquals(5L, result.version()); + assertEquals(42L, result.seqNo()); + assertEquals(2L, result.primaryTerm()); + + String source = result.source().utf8ToString(); + assertTrue("user fields present: " + source, source.contains("\"name\":\"alice\"")); + assertTrue("user fields present: " + source, source.contains("\"age\":30")); + assertFalse("reserved fields excluded: " + source, source.contains("_seq_no")); + assertFalse("reserved fields excluded: " + source, source.contains("_version")); + assertFalse("underscore-prefixed fields excluded: " + source, source.contains("_custom")); + } + + public void testGetById_throwsISEWhenFileSetNull() throws Exception { + when(resolver.resolveMetadata(reader, "doc1")).thenReturn(metadata("doc1", 0L, 7L)); + when(snapshot.findFileSet(FORMAT, 7L)).thenReturn(null); + + IllegalStateException e = expectThrows(IllegalStateException.class, () -> service.getById("doc1", reader, INDEX)); + assertTrue(e.getMessage(), e.getMessage().contains("no matching file set")); + assertTrue(e.getMessage(), e.getMessage().contains("doc1")); + } + + public void testGetById_throwsISEWhenRowNull() throws Exception { + when(resolver.resolveMetadata(reader, "doc1")).thenReturn(metadata("doc1", 0L, 7L)); + WriterFileSet fs = fileSet(7L, "0.parquet"); + when(snapshot.findFileSet(FORMAT, 7L)).thenReturn(fs); + when(executor.executeSingleRow(0L, fs)).thenReturn(null); + + IllegalStateException e = expectThrows(IllegalStateException.class, () -> service.getById("doc1", reader, INDEX)); + assertTrue(e.getMessage(), e.getMessage().contains("backend returned no row")); + assertTrue(e.getMessage(), e.getMessage().contains("doc1")); + } + + // ---- getVersionMetadata ------------------------------------------------- + + public void testGetVersionMetadata_notFoundWhenResolverReturnsNull() throws Exception { + when(resolver.resolveMetadata(reader, "missing")).thenReturn(null); + + DocumentLookupResult result = service.getVersionMetadata("missing", reader, INDEX); + + assertFalse(result.exists()); + assertEquals(Versions.NOT_FOUND, result.version()); + assertNull(result.source()); + } + + public void testGetVersionMetadata_extractsVersionFieldsWithNullSource() throws Exception { + when(resolver.resolveMetadata(reader, "doc1")).thenReturn(metadata("doc1", 0L, 7L)); + WriterFileSet fs = fileSet(7L, "0.parquet"); + when(snapshot.findFileSet(FORMAT, 7L)).thenReturn(fs); + when(executor.executeSingleRow(0L, fs)).thenReturn(row("_seq_no", 42L, "_primary_term", 2L, "_version", 5L, "name", "alice")); + + DocumentLookupResult result = service.getVersionMetadata("doc1", reader, INDEX); + + assertTrue(result.exists()); + assertEquals(5L, result.version()); + assertEquals(42L, result.seqNo()); + assertEquals(2L, result.primaryTerm()); + // version-only hot path: source is never reconstructed + assertNull(result.source()); + assertTrue(result.documentFields().isEmpty()); + } + + public void testGetVersionMetadata_defaultsWhenFieldsMissing() throws Exception { + when(resolver.resolveMetadata(reader, "doc1")).thenReturn(metadata("doc1", 0L, 7L)); + WriterFileSet fs = fileSet(7L, "0.parquet"); + when(snapshot.findFileSet(FORMAT, 7L)).thenReturn(fs); + // row exists (found) but carries no version fields + when(executor.executeSingleRow(0L, fs)).thenReturn(row("name", "alice")); + + DocumentLookupResult result = service.getVersionMetadata("doc1", reader, INDEX); + + assertTrue(result.exists()); + assertEquals(Versions.NOT_FOUND, result.version()); + assertEquals(SequenceNumbers.UNASSIGNED_SEQ_NO, result.seqNo()); + assertEquals(SequenceNumbers.UNASSIGNED_PRIMARY_TERM, result.primaryTerm()); + } + + // ---- getDocsAboveSeqNo -------------------------------------------------- + + public void testGetDocsAboveSeqNo_buildsResultsOnlyForRowsWithId() throws Exception { + Segment segment = Segment.builder(1L).addSearchableFiles(FORMAT, fileSet(1L, "0.parquet")).build(); + when(snapshot.getSegments()).thenReturn(List.of(segment)); + + Map withId = row("_id", "d1", "_seq_no", 10L, "name", "alice"); + Map withoutId = row("_seq_no", 11L, "name", "bob"); + when(executor.executeRowsAboveSeqNo(anyList(), eq(5L))).thenReturn(List.of(withId, withoutId)); + + List results = service.getDocsAboveSeqNo(5L, reader, INDEX); + + assertEquals(1, results.size()); + assertEquals("d1", results.get(0).id()); + assertEquals(10L, results.get(0).seqNo()); + assertTrue(results.get(0).source().utf8ToString().contains("\"name\":\"alice\"")); + } + + public void testGetDocsAboveSeqNo_skipsEmptyAndWrongFormatFileSets() throws Exception { + WriterFileSet good = fileSet(1L, "0.parquet"); + WriterFileSet empty = fileSet(2L); // no files -> excluded + WriterFileSet lucene = fileSet(3L, "seg.cfs"); // different format -> excluded + Segment s1 = Segment.builder(1L).addSearchableFiles(FORMAT, good).build(); + Segment s2 = Segment.builder(2L).addSearchableFiles(FORMAT, empty).build(); + Segment s3 = Segment.builder(3L).addSearchableFiles("lucene", lucene).build(); + when(snapshot.getSegments()).thenReturn(List.of(s1, s2, s3)); + + @SuppressWarnings("unchecked") + ArgumentCaptor> captor = ArgumentCaptor.forClass(List.class); + when(executor.executeRowsAboveSeqNo(captor.capture(), eq(0L))).thenReturn(List.of()); + + service.getDocsAboveSeqNo(0L, reader, INDEX); + + List passed = captor.getValue(); + assertEquals(1, passed.size()); + assertEquals(good, passed.get(0)); + } + + public void testGetDocsAboveSeqNo_emptyWhenNoRows() throws Exception { + when(snapshot.getSegments()).thenReturn(List.of()); + when(executor.executeRowsAboveSeqNo(anyList(), anyLong())).thenReturn(List.of()); + + assertTrue(service.getDocsAboveSeqNo(5L, reader, INDEX).isEmpty()); + } + + // ---- extractLong -------------------------------------------------------- + + public void testExtractLong_number() { + assertEquals(5L, DocumentLookupService.extractLong(Map.of("k", 5L), "k", 99L)); + assertEquals(7L, DocumentLookupService.extractLong(Map.of("k", 7), "k", 99L)); + } + + public void testExtractLong_nullReturnsFallback() { + assertEquals(99L, DocumentLookupService.extractLong(Map.of(), "k", 99L)); + } + + public void testExtractLong_parsesStringNumber() { + assertEquals(123L, DocumentLookupService.extractLong(Map.of("k", "123"), "k", 99L)); + } + + public void testExtractLong_unparseableStringReturnsFallback() { + assertEquals(99L, DocumentLookupService.extractLong(Map.of("k", "abc"), "k", 99L)); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java index beed99f9d0e35..dba4ab4c17990 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java @@ -33,15 +33,20 @@ import org.opensearch.core.common.breaker.CircuitBreaker; import org.opensearch.core.common.io.stream.NamedWriteableRegistry; import org.opensearch.core.common.unit.ByteSizeValue; +import org.opensearch.core.index.Index; import org.opensearch.core.indices.breaker.CircuitBreakerStats; import org.opensearch.core.xcontent.NamedXContentRegistry; import org.opensearch.env.Environment; import org.opensearch.env.NodeEnvironment; import org.opensearch.index.IndexSettings; import org.opensearch.index.IndexSortConfig; +import org.opensearch.index.engine.Engine; import org.opensearch.index.engine.dataformat.DataFormatRegistry; import org.opensearch.index.engine.dataformat.ReaderManagerConfig; +import org.opensearch.index.engine.exec.DocumentMetadataResolver; import org.opensearch.index.engine.exec.EngineReaderManager; +import org.opensearch.index.engine.exec.IndexReaderProvider; +import org.opensearch.index.get.DocumentLookupResult; import org.opensearch.indices.breaker.BreakerSettings; import org.opensearch.monitor.os.OsProbe; import org.opensearch.nativebridge.spi.NativeMemoryFetcher; @@ -50,6 +55,7 @@ import org.opensearch.plugin.stats.AnalyticsBackendTaskCancellationStats; import org.opensearch.plugins.ActionPlugin; import org.opensearch.plugins.CircuitBreakerPlugin; +import org.opensearch.plugins.DocumentLookupProvider; import org.opensearch.plugins.NativeStoreHandle; import org.opensearch.plugins.Plugin; import org.opensearch.plugins.SearchBackEndPlugin; @@ -88,7 +94,8 @@ public class DataFusionPlugin extends Plugin SearchBackEndPlugin, AnalyticsSearchBackendPlugin, ActionPlugin, - CircuitBreakerPlugin { + CircuitBreakerPlugin, + DocumentLookupProvider { private static final Logger logger = LogManager.getLogger(DataFusionPlugin.class); @@ -366,6 +373,8 @@ static String deriveSpillLimitDefault() { private volatile SimpleExtension.ExtensionCollection substraitExtensions; private volatile ClusterService clusterService; private volatile DatafusionSettings datafusionSettings; + // DocumentLookupProvider implementation. Construction deferred until the DataFusion service is live. + private volatile GetService getService; private volatile CircuitBreaker datafusionBreaker; /** @@ -438,6 +447,10 @@ public Collection createComponents( dataFusionService.start(); logger.debug("DataFusion plugin initialized — memory pool {}B, spill limit {}B", memoryPoolLimit, spillMemoryLimit); + // Build the get-by-id service now that the DataFusion runtime is live. The + // DocumentMetadataResolver is supplied per-call by the engine, so it is not needed here. + this.getService = new GetService(this); + // Wire the dynamic spill limit setting to the native runtime so updates via the // cluster settings API take effect without restarting the node. clusterService.getClusterSettings().addSettingsUpdateConsumer(DATAFUSION_SPILL_MEMORY_LIMIT, this::updateSpillMemoryLimit); @@ -805,6 +818,9 @@ public Supplier getAnalyticsBackendNativeMemo @Override public void close() throws IOException { + if (getService != null) { + getService.close(); + } if (dataFusionService != null) { dataFusionService.close(); } @@ -832,4 +848,48 @@ public Map getTopQueriesByMemory() { } return result; } + + /** + * Get-by-id entry point. Delegates to {@link GetService}, which fetches the row through the native runtime. + */ + @Override + public DocumentLookupResult getById(Engine.Get get, IndexReaderProvider.Reader reader, Index index, DocumentMetadataResolver resolver) + throws IOException { + GetService getService = getServiceOrThrow(); + return getService.documentLookupService(resolver).getById(get.id(), reader, index); + } + + @Override + public DocumentLookupResult getVersionMetadata( + String id, + IndexReaderProvider.Reader reader, + Index index, + DocumentMetadataResolver resolver + ) throws IOException { + GetService svc = getServiceOrThrow(); + return svc.documentLookupService(resolver).getVersionMetadata(id, reader, index); + } + + @Override + public List getDocsAboveSeqNo( + long fromSeqNoExclusive, + IndexReaderProvider.Reader reader, + Index index, + DocumentMetadataResolver resolver + ) throws IOException { + GetService svc = getService; + if (svc == null) return List.of(); + return svc.documentLookupService(resolver).getDocsAboveSeqNo(fromSeqNoExclusive, reader, index); + } + + /** + * Returns the {@link GetService} , throwing IllegalStateException if not initialized. + */ + private GetService getServiceOrThrow() { + GetService svc = getService; + if (svc == null) { + throw new IllegalStateException("GetService is not initialized. "); + } + return svc; + } } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/GetService.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/GetService.java new file mode 100644 index 0000000000000..2757d3d58c4df --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/GetService.java @@ -0,0 +1,245 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.be.datafusion; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.analytics.exec.ArrowValues; +import org.opensearch.analytics.spi.DocumentLookupService; +import org.opensearch.analytics.spi.DocumentRowReader; +import org.opensearch.be.datafusion.nativelib.NativeBridge; +import org.opensearch.be.datafusion.nativelib.ReaderHandle; +import org.opensearch.be.datafusion.nativelib.StreamHandle; +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.core.action.ActionListener; +import org.opensearch.index.engine.exec.DocumentMetadataResolver; +import org.opensearch.index.engine.exec.MonoFileWriterSet; +import org.opensearch.index.engine.exec.WriterFileSet; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; +import org.opensearch.index.mapper.IdFieldMapper; +import org.opensearch.index.mapper.Uid; + +import java.io.Closeable; +import java.io.IOException; +import java.lang.foreign.Arena; +import java.lang.foreign.MemorySegment; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; + +/** + * DataFusion-backed get-by-id executor. The resolver maps an {@code _id} to a + * {@code (writerGeneration, rowId)}; this locates the parquet file in the {@link CatalogSnapshot} + * and reads that row via a Substrait scan on the native runtime, flattening the Arrow batch into a + * source map. + */ +@ExperimentalApi +public class GetService implements Closeable { + + private static final Logger logger = LogManager.getLogger(GetService.class); + + private final DocumentRowReader executor; + + /** Production constructor. */ + public GetService(DataFusionPlugin dfPlugin) { + this(new NativeBridgeExecutor(dfPlugin)); + } + + GetService(DocumentRowReader executor) { + this.executor = executor; + } + + /** Returns a core-layer DocumentLookupService wired with this backend's executor and the supplied document resolver. */ + public DocumentLookupService documentLookupService(DocumentMetadataResolver resolver) { + return new DocumentLookupService(resolver, executor); + } + + @Override + public void close() throws IOException { + if (executor instanceof Closeable c) { + c.close(); + } + } + + /** + * Production executor driving {@link NativeBridge}. Spins up a short-lived + * reader scoped to one parquet file, executes the plan, imports the single resulting batch + * via the Arrow C Data Interface, and flattens the first row into a Java map. + */ + static final class NativeBridgeExecutor implements DocumentRowReader, Closeable { + + private static final String GET_BY_ID_TABLE_ALIAS = "_t"; + private static final String PARQUET_FORMAT = "parquet"; + + private final DataFusionPlugin dfPlugin; + private final BufferAllocator sharedAllocator = new RootAllocator(64 * 1024 * 1024); + + NativeBridgeExecutor(DataFusionPlugin dfPlugin) { + this.dfPlugin = dfPlugin; + } + + @Override + public void close() { + sharedAllocator.close(); + } + + @Override + public String formatName() { + return PARQUET_FORMAT; + } + + @Override + public Map executeSingleRow(long rowId, WriterFileSet parquetSet) throws IOException { + if (rowId < 0) { + throw new IllegalArgumentException("rowId must be non-negative, got: " + rowId); + } + String parquetDir = parquetSet.directory(); + String parquetFile = parquetSet.files().iterator().next(); + long runtimePtr = dfPlugin.getDataFusionService().getNativeRuntime().get(); + // ReaderHandle registers the native pointer with NativeHandle so downstream + // validatePointer() calls in executeQueryAsync() find it in the live set. + MonoFileWriterSet segment = MonoFileWriterSet.of(parquetDir, parquetSet.writerGeneration(), parquetFile, 0L); + try (ReaderHandle readerHandle = new ReaderHandle(parquetDir, List.of(segment), null, List.of(), List.of())) { + long readerPtr = readerHandle.getPointer(); + // TODO: only the OFFSET (rowId) varies between calls — cache the Substrait plan + // per reader and rebind the offset to avoid re-planning on every get-by-id. + String sql = "SELECT * FROM \"" + GET_BY_ID_TABLE_ALIAS + "\" LIMIT 1 OFFSET " + Long.toUnsignedString(rowId); + byte[] substraitPlan = NativeBridge.sqlToSubstrait(readerPtr, GET_BY_ID_TABLE_ALIAS, sql, runtimePtr); + WireConfigSnapshot configSnapshot = WireConfigSnapshot.builder(dfPlugin.getDatafusionSettings().getSnapshot()) + .queryStrategy(1) // ListingTable — bypass indexed executor routing + .build(); + long streamPtr = executeNativeQuery( + readerPtr, + substraitPlan, + runtimePtr, + configSnapshot, + "DataFusion get-by-id query failed" + ); + return readSingleRow(streamPtr); + } + } + + @Override + public List> executeRowsAboveSeqNo(List fileSets, long seqNoFloor) throws IOException { + long runtimePtr = dfPlugin.getDataFusionService().getNativeRuntime().get(); + List> all = new ArrayList<>(); + for (WriterFileSet parquetSet : fileSets) { + String parquetDir = parquetSet.directory(); + String parquetFile = parquetSet.files().iterator().next(); + MonoFileWriterSet writerSet = MonoFileWriterSet.of(parquetDir, parquetSet.writerGeneration(), parquetFile, 0L); + try (ReaderHandle readerHandle = new ReaderHandle(parquetDir, List.of(writerSet), null, List.of(), List.of())) { + long readerPtr = readerHandle.getPointer(); + // TODO: only the _seq_no floor varies between calls -- cache the Substrait plan + // per reader and rebind the bound to avoid re-planning on every restore scan. + String sql = "SELECT \"_id\", \"_seq_no\", \"_primary_term\", \"_version\" FROM \"" + + GET_BY_ID_TABLE_ALIAS + + "\" WHERE \"_seq_no\" > " + + seqNoFloor; + byte[] substraitPlan = NativeBridge.sqlToSubstrait(readerPtr, GET_BY_ID_TABLE_ALIAS, sql, runtimePtr); + WireConfigSnapshot configSnapshot = WireConfigSnapshot.builder(dfPlugin.getDatafusionSettings().getSnapshot()) + .queryStrategy(1) // ListingTable — bypass indexed executor routing + .build(); + long streamPtr = executeNativeQuery( + readerPtr, + substraitPlan, + runtimePtr, + configSnapshot, + "DataFusion range query failed" + ); + all.addAll(readAllRows(streamPtr)); + } + } + return all; + } + + private List> readAllRows(long streamPtr) { + List> results = new ArrayList<>(); + try ( + StreamHandle streamHandle = new StreamHandle(streamPtr, dfPlugin.getDataFusionService().getNativeRuntime()); + DatafusionResultStream stream = new DatafusionResultStream(streamHandle, sharedAllocator) + ) { + var iter = stream.iterator(); + while (iter.hasNext()) { + var batch = iter.next(); + try (VectorSchemaRoot root = batch.getArrowRoot()) { + FieldVector idVec = root.getVector(IdFieldMapper.NAME); + for (int i = 0; i < root.getRowCount(); i++) { + Map row = ArrowValues.toSourceMap(root, i); + if (idVec != null && !idVec.isNull(i)) { + row.put(IdFieldMapper.NAME, Uid.decodeId((byte[]) idVec.getObject(i))); + } + results.add(row); + } + } + } + } + return results; + } + + private Map readSingleRow(long streamPtr) { + try ( + StreamHandle streamHandle = new StreamHandle(streamPtr, dfPlugin.getDataFusionService().getNativeRuntime()); + DatafusionResultStream stream = new DatafusionResultStream(streamHandle, sharedAllocator) + ) { + var iter = stream.iterator(); + if (!iter.hasNext()) return null; + var batch = iter.next(); + try (VectorSchemaRoot root = batch.getArrowRoot()) { + if (root.getRowCount() == 0) return null; + return ArrowValues.toSourceMap(root, 0); + } + } + } + + private long executeNativeQuery( + long readerPtr, + byte[] substraitPlan, + long runtimePtr, + WireConfigSnapshot configSnapshot, + String errorMessage + ) throws IOException { + CompletableFuture future = new CompletableFuture<>(); + try (Arena arena = Arena.ofConfined()) { + MemorySegment configSegment = arena.allocate(WireConfigSnapshot.BYTE_SIZE); + configSnapshot.writeTo(configSegment); + NativeBridge.executeQueryAsync( + readerPtr, + GET_BY_ID_TABLE_ALIAS, + substraitPlan, + runtimePtr, + 0L, + configSegment.address(), + new ActionListener<>() { + @Override + public void onResponse(Long v) { + future.complete(v); + } + + @Override + public void onFailure(Exception e) { + future.completeExceptionally(e); + } + } + ); + try { + return future.join(); + } catch (Exception e) { + throw new IOException(errorMessage, e); + } + } + } + + } + +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/GetServiceTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/GetServiceTests.java new file mode 100644 index 0000000000000..8d3fc606ae96a --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/GetServiceTests.java @@ -0,0 +1,155 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.be.datafusion; + +import org.opensearch.analytics.spi.DocumentLookupService; +import org.opensearch.analytics.spi.DocumentRowReader; +import org.opensearch.analytics.spi.ReservedFields; +import org.opensearch.core.index.Index; +import org.opensearch.index.engine.exec.DocumentMetadataResolver; +import org.opensearch.index.engine.exec.DocumentMetadataResolver.DocumentMetadata; +import org.opensearch.index.engine.exec.IndexReaderProvider; +import org.opensearch.index.engine.exec.WriterFileSet; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; +import org.opensearch.index.get.DocumentLookupResult; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; +import java.util.Map; +import java.util.Set; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class GetServiceTests extends OpenSearchTestCase { + + private static final Index INDEX = new Index("idx", "uuid"); + + private DocumentRowReader mockExecutor; + + @Override + public void setUp() throws Exception { + super.setUp(); + mockExecutor = mock(DocumentRowReader.class); + } + + public void testGetById_nullResolver_returnsNotFound() throws IOException { + DocumentMetadataResolver resolver = mock(DocumentMetadataResolver.class); + when(resolver.resolveMetadata(any(), eq("doc1"))).thenReturn(null); + + IndexReaderProvider.Reader mockReader = mock(IndexReaderProvider.Reader.class); + DocumentLookupService service = new GetService(mockExecutor).documentLookupService(resolver); + DocumentLookupResult result = service.getById("doc1", mockReader, INDEX); + assertFalse(result.exists()); + assertEquals("doc1", result.id()); + } + + public void testGetById_docFound_returnsFromParquet() throws IOException { + DocumentMetadataResolver resolver = mock(DocumentMetadataResolver.class); + DocumentMetadata metadata = new DocumentMetadata("doc1", 0L, 1L); + when(resolver.resolveMetadata(any(), eq("doc1"))).thenReturn(metadata); + + DocumentRowReader executor = mock(DocumentRowReader.class); + IndexReaderProvider.Reader mockReader = mock(IndexReaderProvider.Reader.class); + WriterFileSet pset = new WriterFileSet("/dir", 1L, Set.of("f.parquet"), 1L, 100L); + CatalogSnapshot snapshot = mock(CatalogSnapshot.class); + when(mockReader.catalogSnapshot()).thenReturn(snapshot); + when(snapshot.findFileSet("parquet", 1L)).thenReturn(pset); + when(executor.formatName()).thenReturn("parquet"); + when(executor.executeSingleRow(0L, pset)).thenReturn( + Map.of("_id", "doc1", "_seq_no", 5L, "_primary_term", 1L, "_version", 2L, "field1", "value1") + ); + + DocumentLookupService service = new GetService(executor).documentLookupService(resolver); + DocumentLookupResult result = service.getById("doc1", mockReader, INDEX); + + assertTrue("doc should be found", result.exists()); + assertEquals("doc1", result.id()); + assertEquals(5L, result.seqNo()); + assertEquals(2L, result.version()); + } + + public void testGetById_executorReturnsNull_throwsISE() throws IOException { + DocumentMetadataResolver resolver = mock(DocumentMetadataResolver.class); + DocumentMetadata metadata = new DocumentMetadata("doc1", 0L, 1L); + when(resolver.resolveMetadata(any(), eq("doc1"))).thenReturn(metadata); + + DocumentRowReader executor = mock(DocumentRowReader.class); + IndexReaderProvider.Reader mockReader = mock(IndexReaderProvider.Reader.class); + WriterFileSet pset = new WriterFileSet("/dir", 1L, Set.of("f.parquet"), 1L, 100L); + CatalogSnapshot snapshot = mock(CatalogSnapshot.class); + when(mockReader.catalogSnapshot()).thenReturn(snapshot); + when(snapshot.findFileSet("parquet", 1L)).thenReturn(pset); + when(executor.formatName()).thenReturn("parquet"); + when(executor.executeSingleRow(0L, pset)).thenReturn(null); + + DocumentLookupService service = new GetService(executor).documentLookupService(resolver); + expectThrows(IllegalStateException.class, () -> service.getById("doc1", mockReader, INDEX)); + } + + public void testGetById_fileSetNull_throwsISE() throws IOException { + DocumentMetadataResolver resolver = mock(DocumentMetadataResolver.class); + DocumentMetadata metadata = new DocumentMetadata("doc1", 0L, 1L); + when(resolver.resolveMetadata(any(), eq("doc1"))).thenReturn(metadata); + + IndexReaderProvider.Reader mockReader = mock(IndexReaderProvider.Reader.class); + CatalogSnapshot snapshot = mock(CatalogSnapshot.class); + when(mockReader.catalogSnapshot()).thenReturn(snapshot); + when(snapshot.findFileSet("parquet", 1L)).thenReturn(null); + when(mockExecutor.formatName()).thenReturn("parquet"); + + DocumentLookupService service = new GetService(mockExecutor).documentLookupService(resolver); + expectThrows(IllegalStateException.class, () -> service.getById("doc1", mockReader, INDEX)); + } + + public void testGetById_resolverThrows_propagates() throws IOException { + DocumentMetadataResolver resolver = mock(DocumentMetadataResolver.class); + when(resolver.resolveMetadata(any(), eq("doc1"))).thenThrow(new IOException("reader closed")); + + IndexReaderProvider.Reader mockReader = mock(IndexReaderProvider.Reader.class); + DocumentLookupService service = new GetService(mockExecutor).documentLookupService(resolver); + expectThrows(IOException.class, () -> service.getById("doc1", mockReader, INDEX)); + } + + public void testExtractLong_number() throws Exception { + assertEquals(42L, invokeExtractLong(Map.of("k", 42), "k", -1L)); + assertEquals(42L, invokeExtractLong(Map.of("k", 42L), "k", -1L)); + assertEquals(42L, invokeExtractLong(Map.of("k", 42.9), "k", -1L)); + } + + public void testExtractLong_string() throws Exception { + assertEquals(99L, invokeExtractLong(Map.of("k", "99"), "k", -1L)); + } + + public void testExtractLong_missing_returnsFallback() throws Exception { + assertEquals(-1L, invokeExtractLong(Map.of(), "k", -1L)); + } + + public void testExtractLong_invalidString_returnsFallback() throws Exception { + assertEquals(-1L, invokeExtractLong(Map.of("k", "abc"), "k", -1L)); + } + + public void testReservedFields() { + Set reserved = ReservedFields.RESERVED; + assertTrue(reserved.contains("_id")); + assertTrue(reserved.contains("_seq_no")); + assertTrue(reserved.contains("_primary_term")); + assertTrue(reserved.contains("_version")); + } + + public void testNoopResolver_returnsNull() throws IOException { + assertNull(DocumentMetadataResolver.NOOP.resolveMetadata(mock(IndexReaderProvider.Reader.class), "any")); + } + + private static long invokeExtractLong(Map row, String key, long fallback) { + return DocumentLookupService.extractLong(row, key, fallback); + } +} diff --git a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneDocumentResolver.java b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneDocumentResolver.java new file mode 100644 index 0000000000000..3755f74622e2f --- /dev/null +++ b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneDocumentResolver.java @@ -0,0 +1,106 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.be.lucene; + +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FilterLeafReader; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.ReaderUtil; +import org.apache.lucene.index.SegmentReader; +import org.apache.lucene.index.SortedNumericDocValues; +import org.apache.lucene.index.Term; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.TermQuery; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.util.BytesRef; +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.index.engine.dataformat.DocumentInput; +import org.opensearch.index.engine.exec.DocumentMetadataResolver; +import org.opensearch.index.engine.exec.IndexReaderProvider; +import org.opensearch.index.mapper.IdFieldMapper; +import org.opensearch.index.mapper.Uid; + +import java.io.IOException; + +/** + * Lucene-backed {@link DocumentMetadataResolver}: an {@code _id} lookup yields the row location + * ({@code rowId} from {@code __row_id__} doc values, {@code writerGeneration} from the segment + * attribute). Version fields are not read here. An empty reader or absent id reports "not found". + * + * @opensearch.experimental + */ +@ExperimentalApi +public final class LuceneDocumentResolver implements DocumentMetadataResolver { + + @Override + public DocumentMetadata resolveMetadata(IndexReaderProvider.Reader reader, String id) throws IOException { + DirectoryReader luceneReader = directoryReader(reader); + if (luceneReader.numDocs() == 0) { + return null; + } + LeafDoc located = locate(luceneReader, id); + if (located == null) { + return null; + } + LeafReader leaf = located.leaf().reader(); + int localDocId = located.localDocId(); + return new DocumentMetadata(id, readRowId(leaf, localDocId), readWriterGeneration(leaf)); + } + + private long readRowId(LeafReader leaf, int localDocId) throws IOException { + return value(leaf.getSortedNumericDocValues(DocumentInput.ROW_ID_FIELD), localDocId); + } + + private long readWriterGeneration(LeafReader leaf) { + LeafReader unwrapped = FilterLeafReader.unwrap(leaf); + if (!(unwrapped instanceof SegmentReader segmentReader)) { + throw new IllegalStateException("Expected SegmentReader leaf, got " + unwrapped.getClass()); + } + String genAttr = segmentReader.getSegmentInfo().info.getAttribute("writer_generation"); + if (genAttr == null) { + throw new IllegalStateException("Leaf segment missing writer_generation attribute"); + } + try { + return Long.parseLong(genAttr); + } catch (NumberFormatException e) { + throw new IllegalStateException("Invalid writer_generation attribute: [" + genAttr + "]", e); + } + } + + private long value(SortedNumericDocValues dv, int docId) throws IOException { + if (dv == null) { + throw new IllegalStateException("Leaf segment missing " + DocumentInput.ROW_ID_FIELD + " doc values"); + } + if (!dv.advanceExact(docId)) { + throw new IllegalStateException("docId " + docId + " has no " + DocumentInput.ROW_ID_FIELD + " value"); + } + return dv.nextValue(); + } + + private DirectoryReader directoryReader(IndexReaderProvider.Reader reader) { + LuceneReader luceneReader = reader.getReader(LucenePlugin.DATA_FORMAT, LuceneReader.class); + return luceneReader.directoryReader(); + } + + private LeafDoc locate(DirectoryReader luceneReader, String id) throws IOException { + BytesRef idBytes = Uid.encodeId(id); + TopDocs topDocs = new IndexSearcher(luceneReader).search(new TermQuery(new Term(IdFieldMapper.NAME, idBytes)), 1); + if (topDocs.scoreDocs.length == 0) { + return null; + } + int docId = topDocs.scoreDocs[0].doc; + int leafOrd = ReaderUtil.subIndex(docId, luceneReader.leaves()); + LeafReaderContext leafCtx = luceneReader.leaves().get(leafOrd); + return new LeafDoc(leafCtx, docId - leafCtx.docBase); + } + + private record LeafDoc(LeafReaderContext leaf, int localDocId) { + } +} diff --git a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LucenePlugin.java b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LucenePlugin.java index 0fe0d1c5c0191..84b70054a85bc 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LucenePlugin.java +++ b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LucenePlugin.java @@ -36,7 +36,9 @@ import org.opensearch.index.engine.dataformat.IndexingEngineConfig; import org.opensearch.index.engine.dataformat.IndexingExecutionEngine; import org.opensearch.index.engine.dataformat.ReaderManagerConfig; +import org.opensearch.index.engine.exec.DocumentMetadataResolver; import org.opensearch.index.engine.exec.EngineReaderManager; +import org.opensearch.index.engine.exec.IndexReaderProvider; import org.opensearch.index.engine.exec.commit.Committer; import org.opensearch.index.engine.exec.commit.CommitterFactory; import org.opensearch.index.store.checksum.LuceneChecksumHandler; @@ -69,9 +71,16 @@ * @opensearch.experimental */ @ExperimentalApi -public class LucenePlugin extends Plugin implements DataFormatPlugin, SearchBackEndPlugin, EnginePlugin, ActionPlugin { +public class LucenePlugin extends Plugin + implements + DataFormatPlugin, + SearchBackEndPlugin, + EnginePlugin, + ActionPlugin, + DocumentMetadataResolver { public static final LuceneDataFormat DATA_FORMAT = new LuceneDataFormat(); + private final LuceneDocumentResolver documentResolver = new LuceneDocumentResolver(); /** Creates a new LucenePlugin. */ public LucenePlugin() {} @@ -195,4 +204,12 @@ public List getRestHandlers( ) { return List.of(new LuceneStatsRestAction(), new LuceneNodeStatsRestAction()); } + + // --- DocumentMetadataResolver --- + + /** {@inheritDoc} Resolves a document id to its row location via Lucene. */ + @Override + public DocumentMetadata resolveMetadata(IndexReaderProvider.Reader reader, String id) throws IOException { + return documentResolver.resolveMetadata(reader, id); + } } diff --git a/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneDocumentResolverTests.java b/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneDocumentResolverTests.java new file mode 100644 index 0000000000000..14dd8fc05b02d --- /dev/null +++ b/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneDocumentResolverTests.java @@ -0,0 +1,175 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.be.lucene; + +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.FilterCodec; +import org.apache.lucene.codecs.SegmentInfoFormat; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.LongPoint; +import org.apache.lucene.document.NumericDocValuesField; +import org.apache.lucene.document.SortedNumericDocValuesField; +import org.apache.lucene.document.StoredField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.SegmentInfo; +import org.apache.lucene.store.ByteBuffersDirectory; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.IOContext; +import org.opensearch.index.engine.dataformat.DataFormat; +import org.opensearch.index.engine.dataformat.DocumentInput; +import org.opensearch.index.engine.exec.DocumentMetadataResolver.DocumentMetadata; +import org.opensearch.index.engine.exec.IndexReaderProvider; +import org.opensearch.index.mapper.IdFieldMapper; +import org.opensearch.index.mapper.SeqNoFieldMapper; +import org.opensearch.index.mapper.Uid; +import org.opensearch.index.mapper.VersionFieldMapper; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class LuceneDocumentResolverTests extends OpenSearchTestCase { + + private final LuceneDocumentResolver resolver = new LuceneDocumentResolver(); + + // --- resolveMetadata --- + + /** + * When the index has no Lucene secondary reader, {@code directoryReader(reader)} fails fast: + * resolution cannot proceed without a Lucene reader, so it throws rather than returning a bogus result. + */ + public void testResolveMetadata_emptyReader_returnsNull() throws IOException { + try (Directory dir = new ByteBuffersDirectory()) { + try (IndexWriter w = new IndexWriter(dir, new IndexWriterConfig())) { + w.commit(); + } + try (DirectoryReader dr = DirectoryReader.open(dir)) { + assertNull(resolver.resolveMetadata(readerReturning(dr), "doc1")); + } + } + } + + public void testResolveMetadata_idNotPresent_returnsNull() throws IOException { + try (Directory dir = new ByteBuffersDirectory()) { + try (IndexWriter w = new IndexWriter(dir, iwc(codecWithGeneration("1")))) { + addDoc(w, "other", 1L, 1L, 1L, 0L, true); + } + try (DirectoryReader dr = DirectoryReader.open(dir)) { + assertNull(resolver.resolveMetadata(readerReturning(dr), "missing")); + } + } + } + + public void testResolveMetadata_found_returnsMetadata() throws IOException { + try (Directory dir = new ByteBuffersDirectory()) { + try (IndexWriter w = new IndexWriter(dir, iwc(codecWithGeneration("5")))) { + addDoc(w, "doc1", 7L, 42L, 3L, 9L, true); + } + try (DirectoryReader dr = DirectoryReader.open(dir)) { + DocumentMetadata md = resolver.resolveMetadata(readerReturning(dr), "doc1"); + assertNotNull(md); + assertEquals("doc1", md.id()); + assertEquals(9L, md.rowId()); + assertEquals(5L, md.writerGeneration()); + } + } + } + + public void testResolveMetadata_missingWriterGeneration_throwsISE() throws IOException { + try (Directory dir = new ByteBuffersDirectory()) { + // default codec → no writer_generation segment attribute + try (IndexWriter w = new IndexWriter(dir, new IndexWriterConfig())) { + addDoc(w, "doc1", 1L, 1L, 1L, 0L, true); + } + try (DirectoryReader dr = DirectoryReader.open(dir)) { + IllegalStateException e = expectThrows( + IllegalStateException.class, + () -> resolver.resolveMetadata(readerReturning(dr), "doc1") + ); + assertTrue(e.getMessage().contains("writer_generation")); + } + } + } + + public void testResolveMetadata_missingRowIdDocValues_throwsISE() throws IOException { + try (Directory dir = new ByteBuffersDirectory()) { + try (IndexWriter w = new IndexWriter(dir, iwc(codecWithGeneration("1")))) { + addDoc(w, "doc1", 1L, 1L, 1L, 0L, false); // no __row_id__ doc values + } + try (DirectoryReader dr = DirectoryReader.open(dir)) { + IllegalStateException e = expectThrows( + IllegalStateException.class, + () -> resolver.resolveMetadata(readerReturning(dr), "doc1") + ); + assertTrue(e.getMessage().contains(DocumentInput.ROW_ID_FIELD)); + } + } + } + + // --- helpers --- + + private static IndexReaderProvider.Reader readerReturning(DirectoryReader dr) { + IndexReaderProvider.Reader reader = mock(IndexReaderProvider.Reader.class); + when(reader.getReader(any(DataFormat.class), eq(LuceneReader.class))).thenReturn(new LuceneReader(dr, java.util.Map.of())); + return reader; + } + + private static IndexWriterConfig iwc(Codec codec) { + IndexWriterConfig c = new IndexWriterConfig(); + c.setCodec(codec); + return c; + } + + private static void addDoc(IndexWriter w, String id, long version, long seqNo, long primaryTerm, long rowId, boolean withRowId) + throws IOException { + Document doc = new Document(); + doc.add(new StringField(IdFieldMapper.NAME, Uid.encodeId(id), Field.Store.NO)); // for _id TermQuery + doc.add(new StoredField(IdFieldMapper.NAME, Uid.encodeId(id))); // for decodeId (restore path) + doc.add(new NumericDocValuesField(VersionFieldMapper.NAME, version)); + doc.add(new LongPoint(SeqNoFieldMapper.NAME, seqNo)); // for LongPoint range scan + doc.add(new NumericDocValuesField(SeqNoFieldMapper.NAME, seqNo)); + doc.add(new NumericDocValuesField(SeqNoFieldMapper.PRIMARY_TERM_NAME, primaryTerm)); + if (withRowId) { + doc.add(new SortedNumericDocValuesField(DocumentInput.ROW_ID_FIELD, rowId)); + } + w.addDocument(doc); + w.commit(); + } + + /** Codec that stamps a {@code writer_generation} segment attribute, mirroring the production writer. */ + private static Codec codecWithGeneration(String gen) { + Codec base = Codec.getDefault(); + return new FilterCodec(base.getName(), base) { + @Override + public SegmentInfoFormat segmentInfoFormat() { + SegmentInfoFormat delegate = base.segmentInfoFormat(); + return new SegmentInfoFormat() { + @Override + public SegmentInfo read(Directory d, String name, byte[] id, IOContext ctx) throws IOException { + return delegate.read(d, name, id, ctx); + } + + @Override + public void write(Directory d, SegmentInfo info, IOContext ctx) throws IOException { + info.putAttribute("writer_generation", gen); + delegate.write(d, info, ctx); + } + }; + } + }; + } +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ArrowValues.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ArrowValues.java index b62d1c7c7a64c..e79af1e694ff8 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ArrowValues.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ArrowValues.java @@ -10,8 +10,10 @@ import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.MapVector; +import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.util.Text; @@ -50,6 +52,76 @@ public final class ArrowValues { private ArrowValues() {} + private static final DateTimeFormatter ISO_FORMATTER = DateTimeFormatter.ISO_INSTANT.withZone(ZoneOffset.UTC); + + /** Converts row {@code rowId} of a {@link VectorSchemaRoot} into a JSON-friendly field map. */ + public static Map toSourceMap(VectorSchemaRoot root, int rowId) { + Map out = new LinkedHashMap<>(); + for (Field field : root.getSchema().getFields()) { + Object converted = toSourceValue(root.getVector(field.getName()), rowId); + if (converted != null) { + out.put(field.getName(), converted); + } + } + return out; + } + + /** + * Reads an Arrow cell as a JSON-friendly scalar: numerics coerced to + * {@code long}/{@code double}, timestamps rendered as ISO-8601 UTC strings. Binary and + * complex (list/struct/decimal) types are not yet supported and return {@code null}. + */ + public static Object toSourceValue(FieldVector vec, int idx) { + if (vec == null || vec.isNull(idx)) return null; + ArrowType type = vec.getField().getType(); + ArrowType.ArrowTypeID id = type.getTypeID(); + switch (id) { + case Binary: + case LargeBinary: + case FixedSizeBinary: + case BinaryView: + return null; + default: + break; + } + Object raw = vec.getObject(idx); + switch (id) { + case Utf8: + case LargeUtf8: + case Utf8View, Date: + return raw == null ? null : raw.toString(); + case Int: + return raw instanceof Number ? ((Number) raw).longValue() : raw; + case FloatingPoint: + return raw instanceof Number ? ((Number) raw).doubleValue() : raw; + case Bool: + return raw; + case Timestamp: + if (raw instanceof Number) { + ArrowType.Timestamp ts = (ArrowType.Timestamp) type; + return ISO_FORMATTER.format(toInstant(((Number) raw).longValue(), ts.getUnit())); + } + return raw == null ? null : raw.toString(); + default: + // TODO type coverage (list, struct, decimal) + return null; + } + } + + private static Instant toInstant(long v, TimeUnit unit) { + switch (unit) { + case SECOND: + return Instant.ofEpochSecond(v); + case MILLISECOND: + return Instant.ofEpochMilli(v); + case MICROSECOND: + return Instant.ofEpochSecond(v / 1_000_000L, (v % 1_000_000L) * 1_000L); + case NANOSECOND: + default: + return Instant.ofEpochSecond(v / 1_000_000_000L, v % 1_000_000_000L); + } + } + public static Object toJavaValue(FieldVector vector, int index) { if (vector.isNull(index)) return null; if (vector instanceof VarCharVector v) { diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/ArrowValuesTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/ArrowValuesTests.java index 77f3b3088bfc7..6914a1fddca93 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/ArrowValuesTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/ArrowValuesTests.java @@ -11,8 +11,12 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.DateMilliVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.TimeMicroVector; import org.apache.arrow.vector.TimeMilliVector; import org.apache.arrow.vector.TimeNanoVector; @@ -21,7 +25,9 @@ import org.apache.arrow.vector.TimeStampMilliVector; import org.apache.arrow.vector.TimeStampNanoVector; import org.apache.arrow.vector.TimeStampSecVector; +import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.impl.UnionListWriter; import org.apache.arrow.vector.types.TimeUnit; @@ -31,6 +37,10 @@ import org.opensearch.test.OpenSearchTestCase; import java.util.List; +import java.util.Map; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; /** Null, BigInt pass-through, VarChar, Date/Time/Timestamp formatting (scalar + list). */ public class ArrowValuesTests extends OpenSearchTestCase { @@ -212,4 +222,110 @@ public void testListOfTimeElementsHasNoEpochPrefix() { assertEquals(List.of("07:40:05"), cell); } } + + // ---- toSourceValue / toSourceMap (moved from GetService.NativeBridgeExecutor) ---- + + public void testToSourceValueNullReturnsNull() { + try (BigIntVector v = new BigIntVector("x", allocator)) { + v.allocateNew(1); + v.setNull(0); + v.setValueCount(1); + assertNull(ArrowValues.toSourceValue(v, 0)); + } + } + + public void testToSourceValueBinaryDropped() { + try (VarBinaryVector v = new VarBinaryVector("b", allocator)) { + v.allocateNew(); + v.setSafe(0, new byte[] { 1, 2, 3 }); + v.setValueCount(1); + assertNull(ArrowValues.toSourceValue(v, 0)); + } + } + + public void testToSourceValueUtf8AsString() { + try (VarCharVector v = new VarCharVector("s", allocator)) { + v.allocateNew(); + v.setSafe(0, "hello".getBytes(java.nio.charset.StandardCharsets.UTF_8)); + v.setValueCount(1); + assertEquals("hello", ArrowValues.toSourceValue(v, 0)); + } + } + + public void testToSourceValueIntAsLong() { + try (IntVector v = new IntVector("n", allocator)) { + v.allocateNew(1); + v.set(0, 42); + v.setValueCount(1); + assertEquals(42L, ArrowValues.toSourceValue(v, 0)); + } + } + + public void testToSourceValueFloatingPointAsDouble() { + try (Float8Vector v = new Float8Vector("d", allocator)) { + v.allocateNew(1); + v.set(0, 3.5d); + v.setValueCount(1); + assertEquals(3.5d, (double) ArrowValues.toSourceValue(v, 0), 0.0d); + } + } + + public void testToSourceValueBoolAsBoolean() { + try (BitVector v = new BitVector("flag", allocator)) { + v.allocateNew(1); + v.set(0, 1); + v.setValueCount(1); + assertEquals(Boolean.TRUE, ArrowValues.toSourceValue(v, 0)); + } + } + + /** Standard timestamp vectors decode to LocalDateTime; toSourceValue renders that via toString() (ISO-8601, 'T' separator, no 'Z'). */ + public void testToSourceValueTimestampLocalDateTimeToString() { + try (TimeStampMilliVector v = new TimeStampMilliVector("ts", allocator)) { + v.allocateNew(1); + v.set(0, 1_780_731_605_000L); + v.setValueCount(1); + assertEquals("2026-06-06T07:40:05", ArrowValues.toSourceValue(v, 0)); + } + } + + /** When a Timestamp cell decodes to a raw Number, toSourceValue formats it as ISO-8601 UTC (with 'Z') across all units. */ + public void testToSourceValueTimestampNumberRendersIsoUtc() { + assertEquals("2026-06-06T07:40:05Z", numericTimestampSource(1_780_731_605L, TimeUnit.SECOND)); + assertEquals("2026-06-06T07:40:05.123Z", numericTimestampSource(1_780_731_605_123L, TimeUnit.MILLISECOND)); + assertEquals("2026-06-06T07:40:05.123456Z", numericTimestampSource(1_780_731_605_123_456L, TimeUnit.MICROSECOND)); + assertEquals("2026-06-06T07:40:05.123456789Z", numericTimestampSource(1_780_731_605_123_456_789L, TimeUnit.NANOSECOND)); + } + + /** Builds a Timestamp-typed vector whose cell decodes to a raw Long, exercising the numeric ISO branch + toInstant. */ + private static Object numericTimestampSource(long raw, TimeUnit unit) { + FieldVector vec = mock(FieldVector.class); + when(vec.isNull(0)).thenReturn(false); + when(vec.getField()).thenReturn(Field.nullable("ts", new ArrowType.Timestamp(unit, null))); + when(vec.getObject(0)).thenReturn(raw); + return ArrowValues.toSourceValue(vec, 0); + } + + public void testToSourceMapDropsNullsAndPreservesFieldOrder() { + VarCharVector name = new VarCharVector("name", allocator); + IntVector age = new IntVector("age", allocator); + VarCharVector missing = new VarCharVector("missing", allocator); + name.allocateNew(); + name.setSafe(0, "alice".getBytes(java.nio.charset.StandardCharsets.UTF_8)); + name.setValueCount(1); + age.allocateNew(1); + age.set(0, 30); + age.setValueCount(1); + missing.allocateNew(); + missing.setNull(0); + missing.setValueCount(1); + // VectorSchemaRoot.of takes ownership; closing root closes the child vectors. + try (VectorSchemaRoot root = VectorSchemaRoot.of(name, age, missing)) { + Map out = ArrowValues.toSourceMap(root, 0); + assertEquals(2, out.size()); + assertEquals("alice", out.get("name")); + assertEquals(30L, out.get("age")); + assertFalse(out.containsKey("missing")); + } + } } diff --git a/sandbox/plugins/composite-engine/build.gradle b/sandbox/plugins/composite-engine/build.gradle index 06c3e157e2647..143a65571fc95 100644 --- a/sandbox/plugins/composite-engine/build.gradle +++ b/sandbox/plugins/composite-engine/build.gradle @@ -59,6 +59,7 @@ dependencies { internalClusterTestImplementation project(':sandbox:plugins:parquet-data-format') internalClusterTestImplementation project(':sandbox:plugins:analytics-backend-lucene') internalClusterTestImplementation project(':sandbox:plugins:analytics-backend-datafusion') + internalClusterTestImplementation project(':sandbox:plugins:analytics-engine') internalClusterTestImplementation project(':sandbox:libs:analytics-framework') internalClusterTestImplementation project(':sandbox:plugins:native-repository-fs') // Netty4 HTTP transport for ITs that use the low-level RestClient (e.g., DataFormatStatsIT, *DataFormatApi*IT). diff --git a/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/DataFormatAwareReadonlyEngineBaseIT.java b/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/DataFormatAwareReadonlyEngineBaseIT.java index 427b2829f5483..4d0cbb3e3ed23 100644 --- a/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/DataFormatAwareReadonlyEngineBaseIT.java +++ b/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/DataFormatAwareReadonlyEngineBaseIT.java @@ -18,7 +18,6 @@ import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.settings.Settings; import org.opensearch.common.util.FeatureFlags; -import org.opensearch.composite.framework.ParquetOnlyDataFormatPlugin; import org.opensearch.core.common.unit.ByteSizeUnit; import org.opensearch.core.common.unit.ByteSizeValue; import org.opensearch.core.xcontent.NamedXContentRegistry; @@ -29,6 +28,7 @@ import org.opensearch.indices.recovery.RecoverySettings; import org.opensearch.indices.replication.common.ReplicationType; import org.opensearch.node.Node; +import org.opensearch.parquet.ParquetDataFormatPlugin; import org.opensearch.plugins.Plugin; import org.opensearch.remotestore.RemoteStoreBaseIntegTestCase; import org.opensearch.remotestore.mocks.MockFsMetadataSupportedRepositoryPlugin; @@ -136,7 +136,7 @@ protected Collection> nodePlugins() { .filter(p -> p != MockFsRepositoryPlugin.class), Stream.>of( ArrowBasePlugin.class, - ParquetOnlyDataFormatPlugin.class, + ParquetDataFormatPlugin.class, CompositeDataFormatPlugin.class, LucenePlugin.class, DataFusionPlugin.class, diff --git a/server/src/main/java/org/opensearch/index/engine/DataFormatAwareEngine.java b/server/src/main/java/org/opensearch/index/engine/DataFormatAwareEngine.java index b18c4242a5ff8..5c0d6ab08a47c 100644 --- a/server/src/main/java/org/opensearch/index/engine/DataFormatAwareEngine.java +++ b/server/src/main/java/org/opensearch/index/engine/DataFormatAwareEngine.java @@ -14,6 +14,7 @@ import org.apache.lucene.index.Term; import org.apache.lucene.search.ReferenceManager; import org.apache.lucene.store.AlreadyClosedException; +import org.apache.lucene.util.BytesRef; import org.opensearch.common.Booleans; import org.opensearch.common.Nullable; import org.opensearch.common.SetOnce; @@ -23,6 +24,8 @@ import org.opensearch.common.lease.Releasable; import org.opensearch.common.logging.Loggers; import org.opensearch.common.lucene.Lucene; +import org.opensearch.common.lucene.uid.Versions; +import org.opensearch.common.metrics.CounterMetric; import org.opensearch.common.queue.DefaultLockableHolder; import org.opensearch.common.queue.LockablePool; import org.opensearch.common.unit.TimeValue; @@ -55,6 +58,8 @@ import org.opensearch.index.engine.dataformat.merge.OneMerge; import org.opensearch.index.engine.exec.CatalogSnapshotLifecycleListener; import org.opensearch.index.engine.exec.CombinedCatalogSnapshotDeletionPolicy; +import org.opensearch.index.engine.exec.DocumentLookupSupport; +import org.opensearch.index.engine.exec.DocumentMetadataResolver; import org.opensearch.index.engine.exec.EngineReaderManager; import org.opensearch.index.engine.exec.FileDeleter; import org.opensearch.index.engine.exec.FilesListener; @@ -68,6 +73,7 @@ import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import org.opensearch.index.engine.exec.coord.CatalogSnapshotManager; import org.opensearch.index.engine.exec.coord.DataformatAwareCatalogSnapshot; +import org.opensearch.index.get.DocumentLookupResult; import org.opensearch.index.mapper.DocumentMapperForType; import org.opensearch.index.mapper.IdFieldMapper; import org.opensearch.index.mapper.ParsedDocument; @@ -91,6 +97,7 @@ import org.opensearch.index.translog.TranslogOperationHelper; import org.opensearch.index.translog.listener.TranslogEventListener; import org.opensearch.indices.pollingingest.PollingIngestStats; +import org.opensearch.plugins.DocumentLookupProvider; import org.opensearch.search.suggest.completion.CompletionStats; import java.io.Closeable; @@ -167,6 +174,12 @@ public class DataFormatAwareEngine implements Indexer { private final IndexingThrottler throttle; private final AtomicInteger throttleRequestCount = new AtomicInteger(); + @Nullable + private final DocumentLookupProvider documentLookupProvider; + private final DocumentMetadataResolver documentMetadataResolver; + // Shared get-by-id flow (parquet lookup + read version-conflict checks), common to all DataFormatAware* engines. + private final DocumentLookupSupport documentLookup; + // Timestamps and seq-no markers private final AtomicLong maxUnsafeAutoIdTimestamp = new AtomicLong(-1); private final AtomicLong maxSeenAutoIdTimestamp = new AtomicLong(-1); @@ -219,6 +232,9 @@ public class DataFormatAwareEngine implements Indexer { // Latch for the current refresh cycle — decremented by any thread that flushes a writer. // Refresh thread awaits this before proceeding with catalog commit. private volatile CountDownLatch activeFlushLatch; + private final LiveVersionMap versionMap; + private final CounterMetric numVersionLookups = new CounterMetric(); + private final CounterMetric numIndexVersionsLookups = new CounterMetric(); /** * System property to enable or disable pluggable dataformat merge operations. @@ -239,6 +255,16 @@ public class DataFormatAwareEngine implements Indexer { * @param engineConfig the engine configuration */ public DataFormatAwareEngine(EngineConfig engineConfig) { + this(engineConfig, null); + } + + /** + * Constructs a DataFormatBasedEngine. + * + * @param engineConfig the engine configuration + * @param documentLookupProvider the optional plugin powering {@link #getById(Engine.Get, java.util.function.BiFunction)} + */ + public DataFormatAwareEngine(EngineConfig engineConfig, @Nullable DocumentLookupProvider documentLookupProvider) { // DataFormatAwareEngine is the writable primary-side engine. Read-only replicas // (segment-rep) and warm-tier shards must use a read-only engine instead — fail // fast so a misconfiguration surfaces before any indexing or recovery. @@ -258,6 +284,12 @@ public DataFormatAwareEngine(EngineConfig engineConfig) { this.shardId = engineConfig.getShardId(); this.store = engineConfig.getStore(); this.throttle = new IndexingThrottler(); + this.documentLookupProvider = documentLookupProvider != null ? documentLookupProvider : engineConfig.getDocumentLookupProvider(); + this.documentMetadataResolver = engineConfig.getDocumentMetadataResolver() != null + ? engineConfig.getDocumentMetadataResolver() + : DocumentMetadataResolver.NOOP; + this.documentLookup = new DocumentLookupSupport(shardId, this.documentLookupProvider, this.documentMetadataResolver); + this.versionMap = new LiveVersionMap(); List refreshListeners = new ArrayList<>(); if (engineConfig.getInternalRefreshListener() != null) { @@ -408,13 +440,13 @@ public DataFormatAwareEngine(EngineConfig engineConfig) { this.indexingStrategyPlanner = new IndexingStrategyPlanner( engineConfig.getIndexSettings(), engineConfig.getShardId(), - new LiveVersionMap(), + this.versionMap, maxUnsafeAutoIdTimestamp::get, () -> 0L, localCheckpointTracker::getProcessedCheckpoint, this::hasBeenProcessedBefore, op -> OpVsEngineDocStatus.OP_NEWER, - (a, b) -> null, + this::resolveDocVersion, this::updateAutoIdTimestamp, (a, b) -> null ); @@ -444,6 +476,12 @@ public DataFormatAwareEngine(EngineConfig engineConfig) { return gen; } ); + + // Restore version map and checkpoint tracker after crash recovery. + if (localCheckpointTracker.getPersistedCheckpoint() < localCheckpointTracker.getMaxSeqNo()) { + restoreVersionMapAndCheckpointTracker(); + } + // Merge failure cleanup: cleans up unreferenced files and acts as a safety net // for refreshLock. The preMergeCommitHook acquires refreshLock on the merge thread; // if the merge fails before applyMergeChanges can release it, this callback ensures @@ -574,9 +612,12 @@ public Engine.IndexResult index(Engine.Index index) throws IOException { || index.origin() == Engine.Operation.Origin.LOCAL_RESET) : "DataFormatAwareEngine only supports PRIMARY, LOCAL_TRANSLOG_RECOVERY, or LOCAL_RESET origins but got: " + index.origin(); final boolean doThrottle = index.origin().isRecovery() == false; - try (ReleasableLock ignored = readLock.acquire()) { + try (ReleasableLock releasableLock = readLock.acquire()) { ensureOpen(); - try (Releasable indexThrottle = doThrottle ? throttle.acquireThrottle() : () -> {}) { + try ( + Releasable ignored = versionMap.acquireLock(index.uid().bytes()); + Releasable indexThrottle = doThrottle ? throttle.acquireThrottle() : () -> {} + ) { lastWriteNanos = index.startTime(); final IndexingStrategy plan; if (index.origin() == Engine.Operation.Origin.PRIMARY) { @@ -707,6 +748,10 @@ private Engine.IndexResult indexIntoEngine(Engine.Index index, IndexingStrategy final Translog.Location location; if (indexResult.getResultType() == Engine.Result.Type.SUCCESS) { location = translogManager.add(new Translog.Index(index, indexResult)); + versionMap.maybePutIndexUnderLock( + index.uid().bytes(), + new IndexVersionValue(location, indexResult.getVersion(), index.seqNo(), index.primaryTerm()) + ); } else if (indexResult.getSeqNo() != UNASSIGNED_SEQ_NO && indexResult.getFailure() != null && !(indexResult.getFailure() instanceof AppendOnlyIndexOperationRetryException)) { @@ -890,6 +935,7 @@ public void refresh(String source) throws EngineException { refreshLock.lock(); try (GatedCloseable catalogSnapshot = catalogSnapshotManager.acquireSnapshot()) { if (store.tryIncRef()) { + versionMap.beforeRefresh(); try { List>> writers = writerPool.checkoutAll(); List existingSegments = catalogSnapshot.get().getSegments(); @@ -1034,10 +1080,15 @@ public void refresh(String source) throws EngineException { } if (refreshed) { lastRefreshedCheckpointListener.updateRefreshedCheckpoint(localCheckpointBeforeRefresh); + versionMap.pruneTombstones( + engineConfig.getThreadPool().relativeTimeInMillis() - engineConfig.getIndexSettings().getGcDeletesInMillis(), + localCheckpointTracker.getProcessedCheckpoint() + ); triggerPossibleMerges(); // trigger merges } } } finally { + versionMap.afterRefresh(refreshed); IOUtils.close(toClose); refreshLock.unlock(); } @@ -1466,7 +1517,7 @@ public long getMaxSeqNoOfUpdatesOrDeletes() { @Override public void advanceMaxSeqNoOfUpdatesOrDeletes(long maxSeqNoOfUpdatesOnPrimary) { - throw new UnsupportedOperationException("updates/deletes not supported"); + this.maxSeqNoOfUpdatesOrDeletes.updateAndGet(curr -> Math.max(curr, maxSeqNoOfUpdatesOnPrimary)); } @Override @@ -1841,6 +1892,78 @@ public byte[] serializeSnapshotToRemoteMetadata(CatalogSnapshot catalogSnapshot) return catalogSnapshotManager.serializeToCommitFormat(catalogSnapshot); } + /** + * Delegates get-by-id to the installed {@link DocumentLookupProvider}, acquiring a + * per-format reader on the current snapshot and passing it to the plugin. + * Throws {@link UnsupportedOperationException} if no plugin is wired. + */ + @Override + public Engine.GetResult getById(Engine.Get get, BiFunction searcherFactory) + throws IOException { + if (documentLookup.isSupported() == false) { + throw new UnsupportedOperationException("getById not supported: no DocumentLookupProvider installed"); + } + try (ReleasableLock ignored = readLock.acquire()) { + ensureOpen(); + if (get.realtime()) { + VersionValue versionValue; + try (Releasable ignore = versionMap.acquireLock(get.uid().bytes())) { + versionValue = getVersionFromMap(get.uid().bytes()); + } + if (versionValue != null) { + if (versionValue.isDelete()) { + return Engine.GetResult.NOT_EXISTS; + } + if (get.versionType().isVersionConflictForReads(versionValue.version, get.version())) { + throw new VersionConflictEngineException( + shardId, + get.id(), + get.versionType().explainConflictForReads(versionValue.version, get.version()) + ); + } + if (get.getIfSeqNo() != SequenceNumbers.UNASSIGNED_SEQ_NO + && (get.getIfSeqNo() != versionValue.seqNo || get.getIfPrimaryTerm() != versionValue.term)) { + throw new VersionConflictEngineException( + shardId, + get.id(), + get.getIfSeqNo(), + get.getIfPrimaryTerm(), + versionValue.seqNo, + versionValue.term + ); + } + if (get.isReadFromTranslog() && versionValue.getLocation() != null) { + try { + Translog.Operation operation = translogManager.readOperation(versionValue.getLocation()); + if (operation != null) { + Translog.Index index = (Translog.Index) operation; + return new DocumentLookupResult( + get.id(), + index.version(), + true, + index.source(), + index.seqNo(), + index.primaryTerm(), + Map.of(), + Map.of() + ).toGetResult(); + } + } catch (IOException e) { + throw new EngineException(shardId, "failed to read operation from translog", e); + } + } + assert versionValue.seqNo >= 0 : versionValue; + } + } + + // Fall through: read from parquet + try (GatedCloseable readerRef = acquireReader()) { + DocumentLookupResult result = documentLookup.lookupFromReader(get, readerRef.get()); + return result.exists() ? result.toGetResult() : Engine.GetResult.NOT_EXISTS; + } + } // readLock + } + /** * DFA callers MUST use {@link #acquireSafeCatalogSnapshot()} — that API avoids the extra * {@code segments_N} disk read required to materialize a Lucene {@link IndexCommit}, and @@ -1968,6 +2091,7 @@ private void closeNoLock(String reason) { assert rwl.isWriteLockedByCurrentThread() || failEngineLock.isHeldByCurrentThread() : "Either the write lock must be held or the engine must be currently failing"; try { + this.versionMap.clear(); // Discard any pending segments not yet picked up by refresh pendingSegments.clear(); // Close any writers queued for deferred close (their files won't reach the catalog) @@ -2018,6 +2142,126 @@ private void closeReaders() throws IOException { } } + protected VersionValue resolveDocVersion(final Engine.Operation op, boolean loadSeqNo) throws IOException { + assert incrementVersionLookup(); + VersionValue versionValue = getVersionFromMap(op.uid().bytes()); + if (versionValue == null) { + if (documentLookupProvider == null) { + return null; + } + assert incrementIndexVersionLookup(); + DocumentLookupResult lookupResult; + try (GatedCloseable readerRef = acquireReader()) { + Reader reader = readerRef.get(); + if (reader.catalogSnapshot().getSegments().isEmpty()) { + return null; + } + lookupResult = documentLookupProvider.getVersionMetadata(op.id(), reader, shardId.getIndex(), documentMetadataResolver); + } + if (lookupResult.version() != Versions.NOT_FOUND) { + versionValue = new IndexVersionValue(null, lookupResult.version(), lookupResult.seqNo(), lookupResult.primaryTerm()); + } + } else { + if (engineConfig.isEnableGcDeletes() + && versionValue.isDelete() + && (engineConfig.getThreadPool().relativeTimeInMillis() + - ((DeleteVersionValue) versionValue).time) > getGcDeletesInMillis()) { + versionValue = null; + } + } + return versionValue; + } + + long getGcDeletesInMillis() { + return engineConfig.getIndexSettings().getGcDeletesInMillis(); + } + + private boolean incrementVersionLookup() { // only used by asserts + numVersionLookups.inc(); + return true; + } + + private boolean incrementIndexVersionLookup() { + numIndexVersionsLookups.inc(); + return true; + } + + private static OpVsEngineDocStatus compareOpToVersionMapOnSeqNo(String id, long seqNo, long primaryTerm, VersionValue versionValue) { + Objects.requireNonNull(versionValue); + if (seqNo > versionValue.seqNo) { + return OpVsEngineDocStatus.OP_NEWER; + } else if (seqNo == versionValue.seqNo) { + assert versionValue.term == primaryTerm : "primary term not matched; id=" + + id + + " seq_no=" + + seqNo + + " op_term=" + + primaryTerm + + " existing_term=" + + versionValue.term; + return OpVsEngineDocStatus.OP_STALE_OR_EQUAL; + } else { + return OpVsEngineDocStatus.OP_STALE_OR_EQUAL; + } + } + + private void restoreVersionMapAndCheckpointTracker() { + try { + final long persistedCheckpoint = localCheckpointTracker.getPersistedCheckpoint(); + if (documentLookupProvider != null) { + try (GatedCloseable readerRef = acquireReader()) { + List docs = documentLookupProvider.getDocsAboveSeqNo( + persistedCheckpoint, + readerRef.get(), + shardId.getIndex(), + documentMetadataResolver + ); + for (DocumentLookupResult doc : docs) { + localCheckpointTracker.markSeqNoAsProcessed(doc.seqNo()); + localCheckpointTracker.markSeqNoAsPersisted(doc.seqNo()); + final BytesRef uid = new Term(IdFieldMapper.NAME, Uid.encodeId(doc.id())).bytes(); + try (Releasable ignored = versionMap.acquireLock(uid)) { + final VersionValue curr = versionMap.getUnderLock(uid); + if (curr == null + || compareOpToVersionMapOnSeqNo( + doc.id(), + doc.seqNo(), + doc.primaryTerm(), + curr + ) == OpVsEngineDocStatus.OP_NEWER) { + versionMap.putIndexUnderLock( + uid, + new IndexVersionValue(null, doc.version(), doc.seqNo(), doc.primaryTerm()) + ); + } + } + } + } + } + } catch (IOException e) { + throw new EngineCreationFailureException( + config().getShardId(), + "failed to restore version map and local checkpoint tracker", + e + ); + } + } + + private VersionValue getVersionFromMap(BytesRef id) { + if (versionMap.isUnsafe()) { + synchronized (versionMap) { + // we are switching from an unsafe map to a safe map. This might happen concurrently + // but we only need to do this once since the last operation per ID is to add to the version + // map so once we pass this point we can safely lookup from the version map. + if (versionMap.isUnsafe()) { + refresh("unsafe_version_map"); + } + versionMap.enforceSafeAccess(); + } + } + return versionMap.getUnderLock(id); + } + private void awaitPendingClose() { try { closedLatch.await(); diff --git a/server/src/main/java/org/opensearch/index/engine/DataFormatAwareNRTReplicationEngine.java b/server/src/main/java/org/opensearch/index/engine/DataFormatAwareNRTReplicationEngine.java index bcc62279dcc32..8fd3060ef7393 100644 --- a/server/src/main/java/org/opensearch/index/engine/DataFormatAwareNRTReplicationEngine.java +++ b/server/src/main/java/org/opensearch/index/engine/DataFormatAwareNRTReplicationEngine.java @@ -35,6 +35,8 @@ import org.opensearch.index.engine.exec.CatalogSnapshotDeletionPolicy; import org.opensearch.index.engine.exec.CatalogSnapshotLifecycleListener; import org.opensearch.index.engine.exec.CommitFileManager; +import org.opensearch.index.engine.exec.DocumentLookupSupport; +import org.opensearch.index.engine.exec.DocumentMetadataResolver; import org.opensearch.index.engine.exec.EngineReaderManager; import org.opensearch.index.engine.exec.FileDeleter; import org.opensearch.index.engine.exec.Indexer; @@ -45,6 +47,7 @@ import org.opensearch.index.engine.exec.commit.IndexStoreProvider; import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import org.opensearch.index.engine.exec.coord.CatalogSnapshotManager; +import org.opensearch.index.get.DocumentLookupResult; import org.opensearch.index.mapper.DocumentMapperForType; import org.opensearch.index.mapper.IdFieldMapper; import org.opensearch.index.mapper.ParsedDocument; @@ -65,6 +68,7 @@ import org.opensearch.index.translog.TranslogOperationHelper; import org.opensearch.index.translog.WriteOnlyTranslogManager; import org.opensearch.indices.pollingingest.PollingIngestStats; +import org.opensearch.plugins.DocumentLookupProvider; import org.opensearch.search.suggest.completion.CompletionStats; import java.io.Closeable; @@ -89,6 +93,7 @@ import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReentrantLock; import java.util.concurrent.locks.ReentrantReadWriteLock; +import java.util.function.BiFunction; import java.util.function.Supplier; /** @@ -120,6 +125,7 @@ public class DataFormatAwareNRTReplicationEngine implements Indexer { private final CatalogSnapshotManager catalogSnapshotManager; private final Committer committer; private final CatalogSnapshotStatsCache statsCache; + private final DocumentLookupSupport documentLookup; private volatile long lastWriteNanos = System.nanoTime(); private final ReentrantReadWriteLock rwl = new ReentrantReadWriteLock(); private final ReleasableLock readLock = new ReleasableLock(rwl.readLock()); @@ -136,6 +142,10 @@ public DataFormatAwareNRTReplicationEngine(EngineConfig engineConfig) { this.engineConfig = engineConfig; this.shardId = engineConfig.getShardId(); this.store = engineConfig.getStore(); + DocumentMetadataResolver resolver = engineConfig.getDocumentMetadataResolver() != null + ? engineConfig.getDocumentMetadataResolver() + : DocumentMetadataResolver.NOOP; + this.documentLookup = new DocumentLookupSupport(shardId, engineConfig.getDocumentLookupProvider(), resolver); store.incRef(); Map> readerManagersRef = null; @@ -397,6 +407,23 @@ public GatedCloseable acquireReader() throws IOException { } } + /** + * Resolves {@code get} against the replicated segment snapshot via the installed + * {@link DocumentLookupProvider}, with read-time version-conflict checks. No live version map, + * so reads go straight to the snapshot. + */ + @Override + public Engine.GetResult getById(Engine.Get get, BiFunction searcherFactory) + throws IOException { + try (ReleasableLock ignored = readLock.acquire()) { + ensureOpen(); + try (GatedCloseable readerRef = acquireReader()) { + DocumentLookupResult result = documentLookup.getById(get, readerRef.get()); + return result.exists() ? result.toGetResult() : Engine.GetResult.NOT_EXISTS; + } + } + } + @Override public boolean refreshNeeded() { return false; diff --git a/server/src/main/java/org/opensearch/index/engine/DataFormatAwareReadOnlyEngine.java b/server/src/main/java/org/opensearch/index/engine/DataFormatAwareReadOnlyEngine.java index f5568a1c7f229..2fac2d217c73a 100644 --- a/server/src/main/java/org/opensearch/index/engine/DataFormatAwareReadOnlyEngine.java +++ b/server/src/main/java/org/opensearch/index/engine/DataFormatAwareReadOnlyEngine.java @@ -31,6 +31,8 @@ import org.opensearch.index.engine.dataformat.ReaderManagerConfig; import org.opensearch.index.engine.exec.CatalogSnapshotDeletionPolicy; import org.opensearch.index.engine.exec.CatalogSnapshotLifecycleListener; +import org.opensearch.index.engine.exec.DocumentLookupSupport; +import org.opensearch.index.engine.exec.DocumentMetadataResolver; import org.opensearch.index.engine.exec.EngineReaderManager; import org.opensearch.index.engine.exec.FileDeleter; import org.opensearch.index.engine.exec.Indexer; @@ -39,6 +41,7 @@ import org.opensearch.index.engine.exec.commit.IndexStoreProvider; import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import org.opensearch.index.engine.exec.coord.CatalogSnapshotManager; +import org.opensearch.index.get.DocumentLookupResult; import org.opensearch.index.mapper.DocumentMapperForType; import org.opensearch.index.mapper.SourceToParse; import org.opensearch.index.merge.MergeStats; @@ -52,6 +55,7 @@ import org.opensearch.index.translog.TranslogManager; import org.opensearch.index.translog.TranslogStats; import org.opensearch.indices.pollingingest.PollingIngestStats; +import org.opensearch.plugins.DocumentLookupProvider; import org.opensearch.search.suggest.completion.CompletionStats; import java.io.Closeable; @@ -68,6 +72,7 @@ import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.locks.ReentrantLock; import java.util.concurrent.locks.ReentrantReadWriteLock; +import java.util.function.BiFunction; import java.util.function.Supplier; /** @@ -119,12 +124,25 @@ public class DataFormatAwareReadOnlyEngine implements Indexer { // Stats cache — populated once at construction (snapshot is permanent for this engine). private final CatalogSnapshotStatsCache statsCache; + private final DocumentLookupSupport documentLookup; + public DataFormatAwareReadOnlyEngine(EngineConfig engineConfig) { + this(engineConfig, null); + } + + public DataFormatAwareReadOnlyEngine(EngineConfig engineConfig, @Nullable DocumentLookupProvider documentLookupProvider) { this.logger = Loggers.getLogger(DataFormatAwareReadOnlyEngine.class, engineConfig.getShardId()); assert engineConfig.isReadOnlyReplica() == false : "DataFormatAwareReadOnlyEngine must only be created for primary shards; shard " + engineConfig.getShardId(); this.engineConfig = engineConfig; this.shardId = engineConfig.getShardId(); + DocumentLookupProvider provider = documentLookupProvider != null + ? documentLookupProvider + : engineConfig.getDocumentLookupProvider(); + DocumentMetadataResolver resolver = engineConfig.getDocumentMetadataResolver() != null + ? engineConfig.getDocumentMetadataResolver() + : DocumentMetadataResolver.NOOP; + this.documentLookup = new DocumentLookupSupport(shardId, provider, resolver); this.store = engineConfig.getStore(); store.incRef(); @@ -305,6 +323,13 @@ public Engine.NoOpResult noOp(Engine.NoOp noOp) throws IOException { throw new UnsupportedOperationException("DataFormatAwareReadOnlyEngine does not support no-ops"); } + @Override + public Engine.GetResult getById(Engine.Get get, BiFunction searcherFactory) + throws IOException { + DocumentLookupResult result = documentLookup.getById(get, reader); + return result.exists() ? result.toGetResult() : Engine.GetResult.NOT_EXISTS; + } + @Override public Engine.Index prepareIndex( DocumentMapperForType docMapper, diff --git a/server/src/main/java/org/opensearch/index/engine/EngineBackedIndexer.java b/server/src/main/java/org/opensearch/index/engine/EngineBackedIndexer.java index b968a7b6db098..e49a6177638ee 100644 --- a/server/src/main/java/org/opensearch/index/engine/EngineBackedIndexer.java +++ b/server/src/main/java/org/opensearch/index/engine/EngineBackedIndexer.java @@ -32,6 +32,7 @@ import java.io.Closeable; import java.io.IOException; import java.util.List; +import java.util.function.BiFunction; /** * An indexer implementation that uses an engine to perform indexing operations. @@ -468,4 +469,11 @@ public GatedCloseable acquireReader() throws IOException { public Engine getEngine() { return engine; } + + /** Engine-backed get-by-id: delegates to {@link Engine#get(Engine.Get, BiFunction)} with {@code searcherFactory}. */ + @Override + public Engine.GetResult getById(Engine.Get get, BiFunction searcherFactory) + throws IOException { + return engine.get(get, searcherFactory); + } } diff --git a/server/src/main/java/org/opensearch/index/engine/EngineConfig.java b/server/src/main/java/org/opensearch/index/engine/EngineConfig.java index 78e319bfafc3b..2ba5a881c31cb 100644 --- a/server/src/main/java/org/opensearch/index/engine/EngineConfig.java +++ b/server/src/main/java/org/opensearch/index/engine/EngineConfig.java @@ -56,6 +56,7 @@ import org.opensearch.index.codec.CodecService; import org.opensearch.index.codec.CodecSettings; import org.opensearch.index.engine.dataformat.DataFormatRegistry; +import org.opensearch.index.engine.exec.DocumentMetadataResolver; import org.opensearch.index.engine.exec.commit.CommitterFactory; import org.opensearch.index.mapper.DocumentMapperForType; import org.opensearch.index.mapper.MapperService; @@ -69,6 +70,7 @@ import org.opensearch.index.translog.TranslogDeletionPolicyFactory; import org.opensearch.index.translog.TranslogFactory; import org.opensearch.indices.IndexingMemoryController; +import org.opensearch.plugins.DocumentLookupProvider; import org.opensearch.threadpool.ThreadPool; import java.util.Collections; @@ -127,6 +129,10 @@ public final class EngineConfig { private final MapperService mapperService; private final CommitterFactory committerFactory; private final Map checksumStrategies; + @Nullable + private final DocumentLookupProvider documentLookupProvider; + @Nullable + private final DocumentMetadataResolver documentMetadataResolver; /** * A supplier of the outstanding retention leases. This is used during merged operations to determine which operations that have been @@ -321,6 +327,8 @@ private EngineConfig(Builder builder) { this.mapperService = builder.mapperService; this.committerFactory = builder.committerFactory; this.checksumStrategies = builder.checksumStrategies; + this.documentLookupProvider = builder.documentLookupProvider; + this.documentMetadataResolver = builder.documentMetadataResolver; } /** @@ -369,7 +377,9 @@ public Builder toBuilder() { .documentMapperForTypeSupplier(this.documentMapperForTypeSupplier) .indexReaderWarmer(this.indexReaderWarmer) .clusterApplierService(this.clusterApplierService) - .mergedSegmentTransferTracker(this.mergedSegmentTransferTracker); + .mergedSegmentTransferTracker(this.mergedSegmentTransferTracker) + .documentLookupProvider(this.documentLookupProvider) + .documentMetadataResolver(this.documentMetadataResolver); } /** @@ -664,6 +674,18 @@ public Map getChecksumStrategies() { return this.checksumStrategies; } + /** Optional {@link DocumentLookupProvider} for the pluggable get-by-id path, or {@code null}. */ + @Nullable + public DocumentLookupProvider getDocumentLookupProvider() { + return this.documentLookupProvider; + } + + /** Optional {@link DocumentMetadataResolver} passed per-call to the provider, or {@code null}. */ + @Nullable + public DocumentMetadataResolver getDocumentMetadataResolver() { + return this.documentMetadataResolver; + } + /** * Builder for EngineConfig class * @@ -706,6 +728,10 @@ public static class Builder { private MapperService mapperService; private CommitterFactory committerFactory; private Map checksumStrategies = Collections.emptyMap(); + @Nullable + private DocumentLookupProvider documentLookupProvider; + @Nullable + private DocumentMetadataResolver documentMetadataResolver; public Builder shardId(ShardId shardId) { this.shardId = shardId; @@ -882,6 +908,16 @@ public Builder checksumStrategies(Map checksumSt return this; } + public Builder documentLookupProvider(@Nullable DocumentLookupProvider documentLookupProvider) { + this.documentLookupProvider = documentLookupProvider; + return this; + } + + public Builder documentMetadataResolver(@Nullable DocumentMetadataResolver documentMetadataResolver) { + this.documentMetadataResolver = documentMetadataResolver; + return this; + } + public EngineConfig build() { return new EngineConfig(this); } diff --git a/server/src/main/java/org/opensearch/index/engine/EngineConfigFactory.java b/server/src/main/java/org/opensearch/index/engine/EngineConfigFactory.java index b9d5be2ed5f2c..e9e0501f55186 100644 --- a/server/src/main/java/org/opensearch/index/engine/EngineConfigFactory.java +++ b/server/src/main/java/org/opensearch/index/engine/EngineConfigFactory.java @@ -29,6 +29,7 @@ import org.opensearch.index.codec.CodecServiceConfig; import org.opensearch.index.codec.CodecServiceFactory; import org.opensearch.index.engine.dataformat.DataFormatRegistry; +import org.opensearch.index.engine.exec.DocumentMetadataResolver; import org.opensearch.index.engine.exec.commit.CommitterFactory; import org.opensearch.index.mapper.DocumentMapperForType; import org.opensearch.index.mapper.MapperService; @@ -39,6 +40,7 @@ import org.opensearch.index.translog.TranslogConfig; import org.opensearch.index.translog.TranslogDeletionPolicyFactory; import org.opensearch.index.translog.TranslogFactory; +import org.opensearch.plugins.DocumentLookupProvider; import org.opensearch.plugins.EnginePlugin; import org.opensearch.plugins.PluginsService; import org.opensearch.threadpool.ThreadPool; @@ -64,6 +66,10 @@ public class EngineConfigFactory { private final TranslogDeletionPolicyFactory translogDeletionPolicyFactory; private final List additionalCodecs; private final CommitterFactory committerFactory; + @Nullable + private final DocumentLookupProvider documentLookupProvider; + @Nullable + private final DocumentMetadataResolver documentMetadataResolver; /** default ctor primarily used for tests without plugins */ public EngineConfigFactory(IndexSettings idxSettings) { @@ -77,8 +83,31 @@ public EngineConfigFactory(PluginsService pluginsService, IndexSettings idxSetti this(pluginsService.filterPlugins(EnginePlugin.class), idxSettings); } - /* private constructor to construct the factory from specific EnginePlugins and IndexSettings */ + /** + * Factory wiring the optional {@link DocumentLookupProvider} and {@link DocumentMetadataResolver} + * (pluggable get-by-id path) into the produced {@link EngineConfig}. + */ + public EngineConfigFactory( + PluginsService pluginsService, + IndexSettings idxSettings, + @Nullable DocumentLookupProvider documentLookupProvider, + @Nullable DocumentMetadataResolver documentMetadataResolver + ) { + this(pluginsService.filterPlugins(EnginePlugin.class), idxSettings, documentLookupProvider, documentMetadataResolver); + } + + /* package-private constructor from specific EnginePlugins and IndexSettings without document-lookup wiring */ EngineConfigFactory(Collection enginePlugins, IndexSettings idxSettings) { + this(enginePlugins, idxSettings, null, null); + } + + /* private constructor to construct the factory from specific EnginePlugins and IndexSettings */ + EngineConfigFactory( + Collection enginePlugins, + IndexSettings idxSettings, + @Nullable DocumentLookupProvider documentLookupProvider, + @Nullable DocumentMetadataResolver documentMetadataResolver + ) { final List codecRegistries = new ArrayList<>(); Optional codecService = Optional.empty(); String codecServiceOverridingPlugin = null; @@ -149,6 +178,8 @@ public EngineConfigFactory(PluginsService pluginsService, IndexSettings idxSetti this.translogDeletionPolicyFactory = translogDeletionPolicyFactory.orElse((idxs, rtls) -> null); this.additionalCodecs = Collections.unmodifiableList(codecRegistries); this.committerFactory = committerFactories.isEmpty() ? null : committerFactories.getFirst(); + this.documentLookupProvider = documentLookupProvider; + this.documentMetadataResolver = documentMetadataResolver; } /** @@ -229,6 +260,8 @@ public EngineConfig newEngineConfig( .mapperService(mapperService) .committerFactory(committerFactory) .checksumStrategies(checksumStrategies) + .documentLookupProvider(documentLookupProvider) + .documentMetadataResolver(documentMetadataResolver) .build(); } diff --git a/server/src/main/java/org/opensearch/index/engine/exec/DocumentLookupSupport.java b/server/src/main/java/org/opensearch/index/engine/exec/DocumentLookupSupport.java new file mode 100644 index 0000000000000..101bf7256ad76 --- /dev/null +++ b/server/src/main/java/org/opensearch/index/engine/exec/DocumentLookupSupport.java @@ -0,0 +1,103 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.index.engine.exec; + +import org.opensearch.common.Nullable; +import org.opensearch.core.index.shard.ShardId; +import org.opensearch.index.engine.Engine; +import org.opensearch.index.engine.VersionConflictEngineException; +import org.opensearch.index.get.DocumentLookupResult; +import org.opensearch.index.seqno.SequenceNumbers; +import org.opensearch.plugins.DocumentLookupProvider; + +import java.io.IOException; + +/** + * Shared get-by-id helper for the {@code DataFormatAware*} engines. Centralizes the pluggable lookup + * (via the installed {@link DocumentLookupProvider}) and the read-time version-conflict checks so + * the engines don't duplicate them. Immutable and thread-safe; one per shard. + * + * @opensearch.internal + */ +public final class DocumentLookupSupport { + + private final ShardId shardId; + @Nullable + private final DocumentLookupProvider provider; + private final DocumentMetadataResolver resolver; + + public DocumentLookupSupport(ShardId shardId, @Nullable DocumentLookupProvider provider, DocumentMetadataResolver resolver) { + this.shardId = shardId; + this.provider = provider; + this.resolver = resolver; + } + + /** Whether a {@link DocumentLookupProvider} is installed (i.e. get-by-id is supported). */ + public boolean isSupported() { + return provider != null; + } + + /** + * Resolves {@code get} against {@code reader} through the installed provider. + * Returns {@link DocumentLookupResult#notFound} when the reader snapshot has no segments. + * + * @throws UnsupportedOperationException if no {@link DocumentLookupProvider} is installed + */ + public DocumentLookupResult lookupFromReader(Engine.Get get, IndexReaderProvider.Reader reader) throws IOException { + if (provider == null) { + throw new UnsupportedOperationException("getById not supported: no DocumentLookupProvider installed"); + } + if (reader.catalogSnapshot().getSegments().isEmpty()) { + return DocumentLookupResult.notFound(get.id()); + } + return provider.getById(get, reader, shardId.getIndex(), resolver); + } + + /** + * Applies read-time version-conflict checks against a resolved {@code result}, mirroring the + * get semantics in {@code InternalEngine}. A no-op when the document does not exist. + * + * @throws VersionConflictEngineException if the requested version or {@code if_seq_no} / + * {@code if_primary_term} preconditions conflict with the resolved document + */ + public void applyReadVersionConflicts(Engine.Get get, DocumentLookupResult result) { + if (result.exists() == false) { + return; + } + if (get.versionType().isVersionConflictForReads(result.version(), get.version())) { + throw new VersionConflictEngineException( + shardId, + get.id(), + get.versionType().explainConflictForReads(result.version(), get.version()) + ); + } + if (get.getIfSeqNo() != SequenceNumbers.UNASSIGNED_SEQ_NO + && (get.getIfSeqNo() != result.seqNo() || get.getIfPrimaryTerm() != result.primaryTerm())) { + throw new VersionConflictEngineException( + shardId, + get.id(), + get.getIfSeqNo(), + get.getIfPrimaryTerm(), + result.seqNo(), + result.primaryTerm() + ); + } + } + + /** + * Convenience: {@link #lookupFromReader} followed by {@link #applyReadVersionConflicts}. + * Used by the read-only and NRT replica engines, which read directly from the segment + * snapshot without a live version map. + */ + public DocumentLookupResult getById(Engine.Get get, IndexReaderProvider.Reader reader) throws IOException { + DocumentLookupResult result = lookupFromReader(get, reader); + applyReadVersionConflicts(get, result); + return result; + } +} diff --git a/server/src/main/java/org/opensearch/index/engine/exec/DocumentMetadataResolver.java b/server/src/main/java/org/opensearch/index/engine/exec/DocumentMetadataResolver.java new file mode 100644 index 0000000000000..60b5d605c4691 --- /dev/null +++ b/server/src/main/java/org/opensearch/index/engine/exec/DocumentMetadataResolver.java @@ -0,0 +1,55 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.index.engine.exec; + +import org.opensearch.common.annotation.ExperimentalApi; + +import java.io.IOException; + +/** + * Resolves a document's physical row location ({@code rowId}, {@code writerGeneration}) from the + * index's secondary structure, used to fetch the row from the primary store. Returns {@code null} + * ("not found") when there is no id present. + * + * @opensearch.experimental + */ +@ExperimentalApi +public interface DocumentMetadataResolver { + + /** Sentinel for {@link DocumentMetadata} location fields that were not populated. */ + long UNSET = -1L; + + /** No-op resolver used when no backend provides one. */ + DocumentMetadataResolver NOOP = new DocumentMetadataResolver() { + @Override + public DocumentMetadata resolveMetadata(IndexReaderProvider.Reader reader, String id) { + return null; + } + }; + + /** + * Physical row location for a document. Version metadata lives in the primary store, not here. + * + * @param id the document id + * @param rowId the row offset within the data file + * @param writerGeneration the writer generation identifying the data file + */ + @ExperimentalApi + record DocumentMetadata(String id, long rowId, long writerGeneration) { + } + + /** + * Resolve row location for a document id. + * + * @param reader the point-in-time reader snapshot + * @param id the document id to resolve + * @return the {@link DocumentMetadata}, or {@code null} if not found + */ + DocumentMetadata resolveMetadata(IndexReaderProvider.Reader reader, String id) throws IOException; +} diff --git a/server/src/main/java/org/opensearch/index/engine/exec/Indexer.java b/server/src/main/java/org/opensearch/index/engine/exec/Indexer.java index 90908b4f20351..4a1209be0faf4 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/Indexer.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/Indexer.java @@ -12,6 +12,7 @@ import org.opensearch.common.Nullable; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.common.concurrent.GatedCloseable; +import org.opensearch.index.engine.Engine; import org.opensearch.index.engine.EngineConfig; import org.opensearch.index.engine.EngineException; import org.opensearch.index.engine.LifecycleAware; @@ -22,6 +23,7 @@ import java.io.Closeable; import java.io.IOException; +import java.util.function.BiFunction; /** * Unified interface for indexing operations in OpenSearch. @@ -199,6 +201,13 @@ Translog.Snapshot newChangesSnapshot(String source, long fromSeqNo, long toSeqNo */ GatedCloseable acquireLastCommittedSnapshot(boolean flushFirst) throws EngineException, IOException; + /** + * Resolves a document by id. Engine-backed indexers delegate to + * {@link Engine#get(Engine.Get, java.util.function.BiFunction)} with {@code searcherFactory}; + * row-store indexers acquire their own reader and ignore the factory. + */ + Engine.GetResult getById(Engine.Get get, BiFunction searcherFactory) throws IOException; + /** * Returns {@code true} if there are merges queued but not yet started. *

diff --git a/server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshot.java b/server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshot.java index 5c41de84116c2..c4f62478945e2 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshot.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshot.java @@ -243,6 +243,24 @@ public Map> getFilesByFormat() { return cached; } + /** + * Finds the {@link WriterFileSet} for the given data format whose {@code writerGeneration} + * matches, or {@code null} if no segment carries a non-empty file set for that format/generation. + * + * @param format the data format name (e.g. {@code "parquet"}) + * @param writerGeneration the writer generation to match + * @return the matching file set, or {@code null} + */ + public WriterFileSet findFileSet(String format, long writerGeneration) { + for (Segment seg : getSegments()) { + WriterFileSet candidate = seg.dfGroupedSearchableFiles().get(format); + if (candidate == null || candidate.files().isEmpty()) continue; + if (candidate.writerGeneration() != writerGeneration) continue; + return candidate; + } + return null; + } + /** * Sets user-defined metadata for this catalog snapshot. * diff --git a/server/src/main/java/org/opensearch/index/get/DocumentLookupResult.java b/server/src/main/java/org/opensearch/index/get/DocumentLookupResult.java new file mode 100644 index 0000000000000..aedda9afe46a9 --- /dev/null +++ b/server/src/main/java/org/opensearch/index/get/DocumentLookupResult.java @@ -0,0 +1,141 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.index.get; + +import org.opensearch.common.Nullable; +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.common.document.DocumentField; +import org.opensearch.common.lucene.uid.Versions; +import org.opensearch.common.lucene.uid.VersionsAndSeqNoResolver.DocIdAndVersion; +import org.opensearch.core.common.bytes.BytesReference; +import org.opensearch.index.engine.Engine; + +import java.util.Map; +import java.util.Objects; + +import static org.opensearch.index.seqno.SequenceNumbers.UNASSIGNED_PRIMARY_TERM; +import static org.opensearch.index.seqno.SequenceNumbers.UNASSIGNED_SEQ_NO; + +/** + * Minimal result of get-by-id lookup. Decouples callers from + * {@link Engine.GetResult}. + * + * @opensearch.experimental + */ +@ExperimentalApi +public record DocumentLookupResult(String id, long version, boolean exists, @Nullable BytesReference source, long seqNo, long primaryTerm, + Map documentFields, Map metadataFields) { + + public static DocumentLookupResult notFound(String id) { + return new DocumentLookupResult( + id, + Versions.NOT_FOUND, + false, + null, + UNASSIGNED_SEQ_NO, + UNASSIGNED_PRIMARY_TERM, + Map.of(), + Map.of() + ); + } + + public DocumentLookupResult( + String id, + long version, + boolean exists, + @Nullable BytesReference source, + long seqNo, + long primaryTerm, + Map documentFields, + Map metadataFields + ) { + this.id = id; + this.version = version; + this.exists = exists; + this.source = source; + this.seqNo = seqNo; + this.primaryTerm = primaryTerm; + this.documentFields = documentFields == null ? Map.of() : documentFields; + this.metadataFields = metadataFields == null ? Map.of() : metadataFields; + } + + /** + * Wraps this lookup as a {@link PreMaterialized} {@link Engine.GetResult} so the get-by-id path + * shares the {@code IndexShard.get(Engine.Get)} return type. The synthesized + * {@link DocIdAndVersion} carries version/seqNo/primaryTerm only; its reader/docId are unset + * and must not be dereferenced. + */ + public Engine.GetResult toGetResult() { + return new PreMaterialized(this); + } + + /** + * {@link Engine.GetResult} carrying a {@link DocumentLookupResult}. {@link ShardGetService} + * detects it via {@code instanceof} and materializes source/fields from the lookup instead of + * the engine stored-fields path. Has no searcher/docId — those accessors throw. + */ + public static final class PreMaterialized extends Engine.GetResult { + private final DocumentLookupResult lookup; + + private PreMaterialized(DocumentLookupResult lookup) { + super(null, new DocIdAndVersion(-1, lookup.version, lookup.seqNo, lookup.primaryTerm, null, 0), false); + this.lookup = lookup; + } + + public DocumentLookupResult lookup() { + return lookup; + } + + @Override + public Engine.Searcher searcher() { + throw new UnsupportedOperationException("PreMaterialized has no searcher"); + } + + @Override + public DocIdAndVersion docIdAndVersion() { + throw new UnsupportedOperationException("PreMaterialized has no docId"); + } + + @Override + public boolean exists() { + return lookup.exists; + } + + @Override + public long version() { + return lookup.version; + } + + @Override + public void close() { + // no searcher to release + } + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o instanceof DocumentLookupResult == false) return false; + DocumentLookupResult other = (DocumentLookupResult) o; + return version == other.version + && exists == other.exists + && seqNo == other.seqNo + && primaryTerm == other.primaryTerm + && Objects.equals(id, other.id) + && Objects.equals(source, other.source) + && Objects.equals(documentFields, other.documentFields) + && Objects.equals(metadataFields, other.metadataFields); + } + + @Override + public int hashCode() { + return Objects.hash(id, version, exists, source, seqNo, primaryTerm, documentFields, metadataFields); + } + +} diff --git a/server/src/main/java/org/opensearch/index/get/ShardGetService.java b/server/src/main/java/org/opensearch/index/get/ShardGetService.java index c4d765c4e0358..6f8bbc92eb026 100644 --- a/server/src/main/java/org/opensearch/index/get/ShardGetService.java +++ b/server/src/main/java/org/opensearch/index/get/ShardGetService.java @@ -236,11 +236,64 @@ private GetResult innerGet( if (get == null || get.exists() == false) { return new GetResult(shardId.getIndexName(), id, UNASSIGNED_SEQ_NO, UNASSIGNED_PRIMARY_TERM, -1, false, null, null, null); } + if (get instanceof DocumentLookupResult.PreMaterialized) { + return buildFromLookup(((DocumentLookupResult.PreMaterialized) get).lookup(), fetchSourceContext); + } // break between having loaded it from translog (so we only have _source), and having a document to load return innerGetLoadFromStoredFields(id, gFields, fetchSourceContext, get, mapperService); } } + /** + * Builds a {@link GetResult} from a pre-materialized {@link DocumentLookupResult}, applying the + * request-level {@link FetchSourceContext} filters. Used by the pluggable get-by-id path. + */ + private GetResult buildFromLookup(DocumentLookupResult lookup, FetchSourceContext fetchSourceContext) { + BytesReference source = lookup.source(); + if (fetchSourceContext.fetchSource() == false) { + source = null; + } else { + source = applySourceFilter(source, fetchSourceContext); + } + return new GetResult( + shardId.getIndexName(), + lookup.id(), + lookup.seqNo(), + lookup.primaryTerm(), + lookup.version(), + lookup.exists(), + source, + lookup.documentFields(), + lookup.metadataFields() + ); + } + + /** + * Applies request-level {@code includes}/{@code excludes} source filtering via + * {@link XContentMapValues}. Returns the source unchanged when no filters + * are set, and {@code null} when the input is {@code null}. + */ + private static BytesReference applySourceFilter(BytesReference source, FetchSourceContext fetchSourceContext) { + if (source == null) { + return null; + } + if (fetchSourceContext.includes().length > 0 || fetchSourceContext.excludes().length > 0) { + Tuple> typeMapTuple = XContentHelper.convertToMap(source, true); + XContentType sourceContentType = typeMapTuple.v1(); + Map sourceAsMap = XContentMapValues.filter( + typeMapTuple.v2(), + fetchSourceContext.includes(), + fetchSourceContext.excludes() + ); + try { + return BytesReference.bytes(MediaTypeRegistry.contentBuilder(sourceContentType).map(sourceAsMap)); + } catch (IOException e) { + throw new OpenSearchException("Failed to apply source includes/excludes filter", e); + } + } + return source; + } + private GetResult innerGetLoadFromStoredFields( String id, String[] storedFields, @@ -369,28 +422,10 @@ private GetResult innerGetLoadFromStoredFields( } } - if (source != null) { - // apply request-level source filtering - if (fetchSourceContext.fetchSource() == false) { - source = null; - } else if (fetchSourceContext.includes().length > 0 || fetchSourceContext.excludes().length > 0) { - Map sourceAsMap; - // TODO: The source might be parsed and available in the sourceLookup but that one uses unordered maps so different. - // Do we care? - Tuple> typeMapTuple = XContentHelper.convertToMap(source, true); - XContentType sourceContentType = typeMapTuple.v1(); - sourceAsMap = typeMapTuple.v2(); - sourceAsMap = XContentMapValues.filter(sourceAsMap, fetchSourceContext.includes(), fetchSourceContext.excludes()); - try { - source = BytesReference.bytes(MediaTypeRegistry.contentBuilder(sourceContentType).map(sourceAsMap)); - } catch (IOException e) { - throw new OpenSearchException("Failed to get id [" + id + "] with includes/excludes set", e); - } - } - } - - if (!fetchSourceContext.fetchSource()) { + if (fetchSourceContext.fetchSource() == false) { source = null; + } else { + source = applySourceFilter(source, fetchSourceContext); } if (source != null && get.isFromTranslog()) { @@ -402,19 +437,8 @@ private GetResult innerGetLoadFromStoredFields( } } - if (source != null && (fetchSourceContext.includes().length > 0 || fetchSourceContext.excludes().length > 0)) { - Map sourceAsMap; - // TODO: The source might parsed and available in the sourceLookup but that one uses unordered maps so different. Do we care? - Tuple> typeMapTuple = XContentHelper.convertToMap(source, true); - XContentType sourceContentType = typeMapTuple.v1(); - sourceAsMap = typeMapTuple.v2(); - sourceAsMap = XContentMapValues.filter(sourceAsMap, fetchSourceContext.includes(), fetchSourceContext.excludes()); - try { - source = BytesReference.bytes(MediaTypeRegistry.contentBuilder(sourceContentType).map(sourceAsMap)); - } catch (IOException e) { - throw new OpenSearchException("Failed to get id [" + id + "] with includes/excludes set", e); - } - } + // Re-apply includes/excludes after translog reapply (which may have re-added excluded fields). + source = applySourceFilter(source, fetchSourceContext); return new GetResult( shardId.getIndexName(), diff --git a/server/src/main/java/org/opensearch/index/shard/IndexShard.java b/server/src/main/java/org/opensearch/index/shard/IndexShard.java index c5cf064562ba2..522acc858a096 100644 --- a/server/src/main/java/org/opensearch/index/shard/IndexShard.java +++ b/server/src/main/java/org/opensearch/index/shard/IndexShard.java @@ -1532,7 +1532,11 @@ public Engine.GetResult get(Engine.Get get) { if (mapper == null) { return GetResult.NOT_EXISTS; } - return applyOnEngine(getIndexer(), engine -> engine.get(get, this::acquireSearcher)); + try { + return getIndexer().getById(get, this::acquireSearcher); + } catch (IOException e) { + throw new OpenSearchException("get-by-id failed for id [" + get.id() + "]", e); + } } /** diff --git a/server/src/main/java/org/opensearch/indices/IndicesService.java b/server/src/main/java/org/opensearch/indices/IndicesService.java index 9f7d43685dfd9..f9793f71bfb2e 100644 --- a/server/src/main/java/org/opensearch/indices/IndicesService.java +++ b/server/src/main/java/org/opensearch/indices/IndicesService.java @@ -126,6 +126,7 @@ import org.opensearch.index.engine.ReadOnlyEngine; import org.opensearch.index.engine.dataformat.DataFormatRegistry; import org.opensearch.index.engine.exec.DataFormatAwareIndexerFactory; +import org.opensearch.index.engine.exec.DocumentMetadataResolver; import org.opensearch.index.engine.exec.EngineBackedIndexerFactory; import org.opensearch.index.engine.exec.IndexerFactory; import org.opensearch.index.fielddata.IndexFieldDataCache; @@ -172,6 +173,7 @@ import org.opensearch.ingest.IngestService; import org.opensearch.node.Node; import org.opensearch.node.remotestore.RemoteStoreNodeAttribute; +import org.opensearch.plugins.DocumentLookupProvider; import org.opensearch.plugins.IndexStorePlugin; import org.opensearch.plugins.PluginsService; import org.opensearch.plugins.SearchStatsContributor; @@ -1234,7 +1236,7 @@ private synchronized IndexService createIndexService( } private EngineConfigFactory getEngineConfigFactory(final IndexSettings idxSettings) { - return new EngineConfigFactory(this.pluginsService, idxSettings); + return new EngineConfigFactory(this.pluginsService, idxSettings, resolveGetByIdPlugin(), resolveDocumentResolver()); } private IngestionConsumerFactory getIngestionConsumerFactory(final IndexSettings idxSettings) { @@ -1260,6 +1262,28 @@ private IndexerFactory getIndexerFactory(final IndexSettings idxSettings) { } } + private DocumentLookupProvider resolveGetByIdPlugin() { + List plugins = pluginsService.filterPlugins(DocumentLookupProvider.class); + if (plugins.isEmpty()) { + return null; + } + if (plugins.size() > 1) { + throw new IllegalStateException("multiple DocumentLookupProvider implementations registered: " + plugins); + } + return plugins.getFirst(); + } + + private DocumentMetadataResolver resolveDocumentResolver() { + List resolvers = pluginsService.filterPlugins(DocumentMetadataResolver.class); + if (resolvers.isEmpty()) { + return DocumentMetadataResolver.NOOP; + } + if (resolvers.size() > 1) { + throw new IllegalStateException("multiple DocumentMetadataResolver implementations registered: " + resolvers); + } + return resolvers.getFirst(); + } + private EngineFactory getEngineFactory(final IndexSettings idxSettings) { final IndexMetadata indexMetadata = idxSettings.getIndexMetadata(); if (indexMetadata != null && indexMetadata.getState() == IndexMetadata.State.CLOSE) { diff --git a/server/src/main/java/org/opensearch/plugins/DocumentLookupProvider.java b/server/src/main/java/org/opensearch/plugins/DocumentLookupProvider.java new file mode 100644 index 0000000000000..61bc2f7bdd99c --- /dev/null +++ b/server/src/main/java/org/opensearch/plugins/DocumentLookupProvider.java @@ -0,0 +1,70 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugins; + +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.core.index.Index; +import org.opensearch.index.engine.Engine; +import org.opensearch.index.engine.exec.DocumentMetadataResolver; +import org.opensearch.index.engine.exec.IndexReaderProvider; +import org.opensearch.index.get.DocumentLookupResult; + +import java.io.IOException; +import java.util.List; + +/** + * SPI for pluggable get-by-id lookup. Implementations resolve a document id + * against the shard's current reader snapshot and return a + * {@link DocumentLookupResult} with the source (and optionally other fields) + * for the matching row. + * + * @opensearch.experimental + */ +@ExperimentalApi +public interface DocumentLookupProvider { + + /** + * Resolve the document referenced by {@code get} against {@code reader}. + * + * @param get the get request (id, realtime flag, etc.) + * @param reader a point-in-time index reader acquired by the caller + * @param index the index, used by implementations as the scan table name + * @param resolver the {@link DocumentMetadataResolver} used to resolve a document's row location + * @return the lookup result; never {@code null} + * @throws IOException if resolution fails + */ + DocumentLookupResult getById(Engine.Get get, IndexReaderProvider.Reader reader, Index index, DocumentMetadataResolver resolver) + throws IOException; + + /** + * Resolves only version metadata ({@code _version}/{@code _seq_no}/{@code _primary_term}) for an id, + * skipping {@code _source} reconstruction. The {@code resolver} resolves the document's row location. + */ + default DocumentLookupResult getVersionMetadata( + String id, + IndexReaderProvider.Reader reader, + Index index, + DocumentMetadataResolver resolver + ) throws IOException { + return DocumentLookupResult.notFound(id); + } + + /** + * Returns metadata for all documents with {@code _seq_no > fromSeqNoExclusive}. Default returns empty list. + * The {@code resolver} resolves each document's row location. + */ + default List getDocsAboveSeqNo( + long fromSeqNoExclusive, + IndexReaderProvider.Reader reader, + Index index, + DocumentMetadataResolver resolver + ) throws IOException { + return List.of(); + } +} diff --git a/server/src/test/java/org/opensearch/index/engine/DataFormatAwareEngineTests.java b/server/src/test/java/org/opensearch/index/engine/DataFormatAwareEngineTests.java index 28a7440828654..dd5a5ee3d2769 100644 --- a/server/src/test/java/org/opensearch/index/engine/DataFormatAwareEngineTests.java +++ b/server/src/test/java/org/opensearch/index/engine/DataFormatAwareEngineTests.java @@ -48,6 +48,7 @@ import org.opensearch.index.engine.exec.commit.Committer; import org.opensearch.index.engine.exec.commit.CommitterFactory; import org.opensearch.index.engine.exec.coord.CatalogSnapshot; +import org.opensearch.index.get.DocumentLookupResult; import org.opensearch.index.mapper.DocumentMapper; import org.opensearch.index.mapper.IdFieldMapper; import org.opensearch.index.mapper.MapperService; @@ -61,6 +62,7 @@ import org.opensearch.index.store.Store; import org.opensearch.index.translog.Translog; import org.opensearch.index.translog.TranslogConfig; +import org.opensearch.plugins.DocumentLookupProvider; import org.opensearch.plugins.PluginsService; import org.opensearch.plugins.SearchBackEndPlugin; import org.opensearch.test.DummyShardLock; @@ -91,6 +93,8 @@ import static org.hamcrest.Matchers.greaterThanOrEqualTo; import static org.hamcrest.Matchers.instanceOf; import static org.hamcrest.Matchers.notNullValue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -159,6 +163,10 @@ private Store createStore() throws IOException { * expects to find: translog UUID, seq-no info, history UUID, etc. */ private void bootstrapStoreWithMetadata(Store store, String translogUUID) throws IOException { + bootstrapStoreWithMetadata(store, translogUUID, SequenceNumbers.NO_OPS_PERFORMED); + } + + private void bootstrapStoreWithMetadata(Store store, String translogUUID, long maxSeqNo) throws IOException { try ( IndexWriter writer = new IndexWriter( store.directory(), @@ -169,7 +177,7 @@ private void bootstrapStoreWithMetadata(Store store, String translogUUID) throws Map commitData = new HashMap<>(); commitData.put(Translog.TRANSLOG_UUID_KEY, translogUUID); commitData.put(SequenceNumbers.LOCAL_CHECKPOINT_KEY, Long.toString(SequenceNumbers.NO_OPS_PERFORMED)); - commitData.put(SequenceNumbers.MAX_SEQ_NO, Long.toString(SequenceNumbers.NO_OPS_PERFORMED)); + commitData.put(SequenceNumbers.MAX_SEQ_NO, Long.toString(maxSeqNo)); commitData.put(Engine.MAX_UNSAFE_AUTO_ID_TIMESTAMP_COMMIT_ID, "-1"); commitData.put(Engine.HISTORY_UUID_KEY, UUID.randomUUID().toString()); writer.setLiveCommitData(commitData.entrySet()); @@ -283,6 +291,64 @@ private Engine.Index indexOp(ParsedDocument doc) { ); } + /** + * Append-only optimizable variant of {@link #indexOp}: a non-negative autoGeneratedIdTimestamp + * with isRetry=false makes {@link IndexingStrategyPlanner} take the optimizedAppendOnly path + * (no safe-access enforcement). + */ + private Engine.Index appendOnlyOp(ParsedDocument doc) { + return new Engine.Index( + new Term(IdFieldMapper.NAME, Uid.encodeId(doc.id())), + doc, + SequenceNumbers.UNASSIGNED_SEQ_NO, + primaryTerm.get(), + Versions.MATCH_ANY, + VersionType.INTERNAL, + Engine.Operation.Origin.PRIMARY, + System.nanoTime(), + System.currentTimeMillis(), + false, + SequenceNumbers.UNASSIGNED_SEQ_NO, + 0 + ); + } + + /** Explicit-id index with a caller-supplied version/versionType (for version-conflict tests). */ + private Engine.Index indexOpWithVersion(ParsedDocument doc, long version, VersionType versionType) { + return new Engine.Index( + new Term(IdFieldMapper.NAME, Uid.encodeId(doc.id())), + doc, + SequenceNumbers.UNASSIGNED_SEQ_NO, + primaryTerm.get(), + version, + versionType, + Engine.Operation.Origin.PRIMARY, + System.nanoTime(), + -1, + false, + SequenceNumbers.UNASSIGNED_SEQ_NO, + 0 + ); + } + + /** Explicit-id index with compare-and-set ifSeqNo/ifPrimaryTerm (for optimistic-concurrency tests). */ + private Engine.Index indexOpWithIfSeqNo(ParsedDocument doc, long ifSeqNo, long ifPrimaryTerm) { + return new Engine.Index( + new Term(IdFieldMapper.NAME, Uid.encodeId(doc.id())), + doc, + SequenceNumbers.UNASSIGNED_SEQ_NO, + primaryTerm.get(), + Versions.MATCH_ANY, + VersionType.INTERNAL, + Engine.Operation.Origin.PRIMARY, + System.nanoTime(), + -1, + false, + ifSeqNo, + ifPrimaryTerm + ); + } + /** * Wraps {@link EngineTestCase#createParsedDoc(String, String)} to attach a * {@link MockDocumentInput}. {@link DataFormatAwareEngine#indexIntoEngine} requires a @@ -3553,4 +3619,440 @@ public void testOnSettingsChangedUnrecognizedTieringValueNotFrozen() throws IOEx assertEquals(Engine.Result.Type.SUCCESS, result.getResultType()); } } + + private DataFormatAwareEngine createDFAEngineWithLookupProvider(Store store, Path translogPath, DocumentLookupProvider provider) + throws IOException { + String uuid = Translog.createEmptyTranslog(translogPath, SequenceNumbers.NO_OPS_PERFORMED, shardId, primaryTerm.get()); + bootstrapStoreWithMetadata(store, uuid); + return new DataFormatAwareEngine(buildDFAEngineConfig(store, translogPath), provider); + } + + private Engine.Get realtimeGet(String id) { + return new Engine.Get(true, true, id, new Term(IdFieldMapper.NAME, Uid.encodeId(id))); + } + + private DocumentLookupProvider mockLookupProvider() throws IOException { + DocumentLookupProvider provider = mock(DocumentLookupProvider.class); + when(provider.getById(any(), any(), any(), any())).thenReturn(DocumentLookupResult.notFound("")); + when(provider.getVersionMetadata(any(), any(), any(), any())).thenReturn(DocumentLookupResult.notFound("")); + return provider; + } + + public void testGetMaxSeqNoOfUpdatesOrDeletes() throws IOException { + try (DataFormatAwareEngine engine = createDFAEngine(store, createTempDir())) { + assertThat(engine.getMaxSeqNoOfUpdatesOrDeletes(), equalTo(SequenceNumbers.NO_OPS_PERFORMED)); + } + } + + public void testCurrentOngoingRefreshCheckpoint() throws IOException { + try (DataFormatAwareEngine engine = createDFAEngine(store, createTempDir())) { + assertThat(engine.currentOngoingRefreshCheckpoint(), equalTo(SequenceNumbers.NO_OPS_PERFORMED)); + for (int i = 0; i < 5; i++) { + engine.index(indexOp(createParsedDocWithInput(Integer.toString(i), null))); + } + engine.refresh("test"); + assertThat(engine.currentOngoingRefreshCheckpoint(), greaterThanOrEqualTo(0L)); + } + } + + /** Calls the new getById(Engine.Get, searcherFactory) and unwraps the pre-materialized lookup. + * DFA engines ignore the searcher factory, so a no-op factory is supplied. */ + private static DocumentLookupResult getByIdLookup(DataFormatAwareEngine engine, Engine.Get get) throws IOException { + Engine.GetResult result = engine.getById(get, (source, scope) -> null); + return result.exists() ? ((DocumentLookupResult.PreMaterialized) result).lookup() : DocumentLookupResult.notFound(get.id()); + } + + public void testGetByIdThrowsWhenNoProvider() throws IOException { + try (DataFormatAwareEngine engine = createDFAEngine(store, createTempDir())) { + Engine.Get get = realtimeGet("1"); + expectThrows(UnsupportedOperationException.class, () -> getByIdLookup(engine, get)); + } + } + + public void testGetByIdReturnsNotFoundWhenEmptyCatalog() throws IOException { + DocumentLookupProvider provider = mockLookupProvider(); + try (DataFormatAwareEngine engine = createDFAEngineWithLookupProvider(store, createTempDir(), provider)) { + Engine.Get get = realtimeGet("1"); + DocumentLookupResult result = getByIdLookup(engine, get); + assertFalse("should be not found on empty catalog", result.exists()); + } + } + + public void testGetByIdReturnsDocFromTranslog() throws IOException { + DocumentLookupProvider provider = mockLookupProvider(); + try (DataFormatAwareEngine engine = createDFAEngineWithLookupProvider(store, createTempDir(), provider)) { + engine.index(indexOp(createParsedDocWithInput("1", null))); + // Realtime get with readFromTranslog finds it in translog via versionMap location + Engine.Get get = realtimeGet("1"); + DocumentLookupResult result = getByIdLookup(engine, get); + assertTrue("doc should be found from translog", result.exists()); + assertThat(result.seqNo(), equalTo(0L)); + } + } + + public void testGetByIdReturnsDocFromParquetAfterRefresh() throws IOException { + DocumentLookupProvider provider = mock(DocumentLookupProvider.class); + when(provider.getById(any(), any(), any(), any())).thenReturn( + new DocumentLookupResult("1", 1L, true, null, 0L, 1L, Map.of(), Map.of()) + ); + try (DataFormatAwareEngine engine = createDFAEngineWithLookupProvider(store, createTempDir(), provider)) { + engine.index(indexOp(createParsedDocWithInput("1", null))); + engine.refresh("test"); + // Non-realtime get falls through to provider + Engine.Get get = new Engine.Get(false, false, "1", new Term(IdFieldMapper.NAME, Uid.encodeId("1"))); + DocumentLookupResult result = getByIdLookup(engine, get); + assertTrue("doc should be found from parquet", result.exists()); + } + } + + /** + * Covers resolveDocVersion falling back to DocumentLookupProvider when versionMap miss. + * Also covers incrementIndexVersionLookup (called inside resolveDocVersion on provider path). + */ + public void testResolveDocVersionFallsBackToProvider() throws IOException { + DocumentLookupProvider provider = mock(DocumentLookupProvider.class); + when(provider.getVersionMetadata(any(), any(), any(), any())).thenReturn( + new DocumentLookupResult("1", 5L, true, null, 3L, 1L, Map.of(), Map.of()) + ); + try (DataFormatAwareEngine engine = createDFAEngineWithLookupProvider(store, createTempDir(), provider)) { + engine.index(indexOp(createParsedDocWithInput("1", null))); + engine.refresh("test"); + // Re-index same doc — resolveDocVersion misses versionMap (cleared by refresh rotation), + // falls back to provider which returns version=5. MATCH_ANY accepts. + Engine.IndexResult result = engine.index(indexOp(createParsedDocWithInput("1", null))); + assertThat(result.getResultType(), equalTo(Engine.Result.Type.SUCCESS)); + assertThat(result.getSeqNo(), equalTo(1L)); + } + } + + /** Reflectively obtains the engine's LiveVersionMap. resolveDocVersion/getVersionFromMap assert the + * per-uid keyed lock is held, so callers must wrap in {@code versionMap.acquireLock(uid)}. */ + @SuppressForbidden(reason = "test needs reflective access to the engine's versionMap field") + private LiveVersionMap versionMapOf(DataFormatAwareEngine engine) throws Exception { + java.lang.reflect.Field vmField = DataFormatAwareEngine.class.getDeclaredField("versionMap"); + vmField.setAccessible(true); + return (LiveVersionMap) vmField.get(engine); + } + + /** resolveDocVersion branch 1: versionMap miss with no provider returns null. */ + public void testResolveDocVersionReturnsNullWhenNoProvider() throws Exception { + try (DataFormatAwareEngine engine = createDFAEngine(store, createTempDir())) { + Engine.Index op = indexOp(createParsedDocWithInput("1", null)); + try (org.opensearch.common.lease.Releasable ignored = versionMapOf(engine).acquireLock(op.uid().bytes())) { + assertNull(engine.resolveDocVersion(op, true)); + } + } + } + + /** resolveDocVersion branch 2: versionMap miss with an empty catalog returns null (provider not consulted). */ + public void testResolveDocVersionReturnsNullOnEmptyCatalog() throws Exception { + try (DataFormatAwareEngine engine = createDFAEngineWithLookupProvider(store, createTempDir(), mockLookupProvider())) { + Engine.Index op = indexOp(createParsedDocWithInput("1", null)); + try (org.opensearch.common.lease.Releasable ignored = versionMapOf(engine).acquireLock(op.uid().bytes())) { + assertNull(engine.resolveDocVersion(op, true)); + } + } + } + + /** resolveDocVersion branch 4: versionMap miss, provider consulted, getById returns NOT_FOUND → null. */ + public void testResolveDocVersionReturnsNullWhenProviderNotFound() throws Exception { + try (DataFormatAwareEngine engine = createDFAEngineWithLookupProvider(store, createTempDir(), mockLookupProvider())) { + engine.index(indexOp(createParsedDocWithInput("1", null))); + engine.refresh("test"); + Engine.Index op = indexOp(createParsedDocWithInput("2", null)); + try (org.opensearch.common.lease.Releasable ignored = versionMapOf(engine).acquireLock(op.uid().bytes())) { + assertNull(engine.resolveDocVersion(op, true)); + } + } + } + + /** resolveDocVersion branch 6: versionMap hit returns the stored IndexVersionValue. */ + public void testResolveDocVersionReturnsVersionMapHit() throws Exception { + try (DataFormatAwareEngine engine = createDFAEngine(store, createTempDir())) { + engine.index(indexOp(createParsedDocWithInput("1", null))); + Engine.Index op = indexOp(createParsedDocWithInput("1", null)); + VersionValue vv; + try (org.opensearch.common.lease.Releasable ignored = versionMapOf(engine).acquireLock(op.uid().bytes())) { + vv = engine.resolveDocVersion(op, true); + } + assertNotNull(vv); + assertThat(vv.version, equalTo(1L)); + assertThat(vv.seqNo, equalTo(0L)); + } + } + + /** + * Covers restoreVersionMapAndCheckpointTracker + compareOpToVersionMapOnSeqNo. + * The provider returns docs for seqNos 0-3 so checkpoint advances contiguously. + */ + public void testRestoreVersionMapAndCheckpointTracker() throws IOException { + Path translogPath = createTempDir(); + DocumentLookupProvider provider = mock(DocumentLookupProvider.class); + when(provider.getDocsAboveSeqNo(anyLong(), any(), any(), any())).thenReturn( + List.of( + new DocumentLookupResult("a", 1L, true, null, 0L, 1L, Map.of(), Map.of()), + new DocumentLookupResult("b", 1L, true, null, 1L, 1L, Map.of(), Map.of()), + new DocumentLookupResult("c", 1L, true, null, 2L, 1L, Map.of(), Map.of()), + new DocumentLookupResult("d", 1L, true, null, 3L, 1L, Map.of(), Map.of()) + ) + ); + when(provider.getVersionMetadata(any(), any(), any(), any())).thenReturn(DocumentLookupResult.notFound("x")); + + // Bootstrap with maxSeqNo=3, localCheckpoint=-1 (triggers restore) + String uuid = Translog.createEmptyTranslog(translogPath, SequenceNumbers.NO_OPS_PERFORMED, shardId, primaryTerm.get()); + bootstrapStoreWithMetadata(store, uuid, 3L); + + EngineConfig config = buildDFAEngineConfig(store, translogPath); + try (DataFormatAwareEngine engine = new DataFormatAwareEngine(config, provider)) { + // All seqNos 0-3 marked → checkpoint advances to 3 + assertThat(engine.getProcessedLocalCheckpoint(), equalTo(3L)); + } + } + + /** + * Covers compareOpToVersionMapOnSeqNo returning OP_STALE_OR_EQUAL. + * Provider returns entries for id "x" at seqNo=2 and seqNo=5, plus a duplicate at seqNo=5. + * The stale entry (seqNo=2) is overwritten by OP_NEWER, the duplicate (same seqNo=5) is skipped. + */ + public void testRestoreVersionMapSkipsStaleEntries() throws IOException { + Path translogPath = createTempDir(); + DocumentLookupProvider provider = mock(DocumentLookupProvider.class); + // Return all seqNos 0-5 so checkpoint advances, with duplicate id "x" at seqNo 2 and 5 + when(provider.getDocsAboveSeqNo(anyLong(), any(), any(), any())).thenReturn( + List.of( + new DocumentLookupResult("a", 1L, true, null, 0L, 1L, Map.of(), Map.of()), + new DocumentLookupResult("b", 1L, true, null, 1L, 1L, Map.of(), Map.of()), + new DocumentLookupResult("x", 1L, true, null, 2L, 1L, Map.of(), Map.of()), + new DocumentLookupResult("c", 1L, true, null, 3L, 1L, Map.of(), Map.of()), + new DocumentLookupResult("d", 1L, true, null, 4L, 1L, Map.of(), Map.of()), + new DocumentLookupResult("x", 2L, true, null, 5L, 1L, Map.of(), Map.of()), + // Duplicate id "x" at same seqNo=5 → hits seqNo == versionValue.seqNo → OP_STALE_OR_EQUAL + new DocumentLookupResult("x", 2L, true, null, 5L, 1L, Map.of(), Map.of()) + ) + ); + when(provider.getVersionMetadata(any(), any(), any(), any())).thenReturn(DocumentLookupResult.notFound("x")); + + String uuid = Translog.createEmptyTranslog(translogPath, SequenceNumbers.NO_OPS_PERFORMED, shardId, primaryTerm.get()); + bootstrapStoreWithMetadata(store, uuid, 5L); + + EngineConfig config = buildDFAEngineConfig(store, translogPath); + try (DataFormatAwareEngine engine = new DataFormatAwareEngine(config, provider)) { + // All seqNos 0-5 contiguous → checkpoint = 5 + assertThat(engine.getProcessedLocalCheckpoint(), equalTo(5L)); + // Index "x" again — should succeed (versionMap has version=2 from seqNo=5 entry) + Engine.IndexResult result = engine.index(indexOp(createParsedDocWithInput("x", null))); + assertThat(result.getResultType(), equalTo(Engine.Result.Type.SUCCESS)); + } + } + + public void testRestoreVersionMapThrowsOnIOException() throws IOException { + Path translogPath = createTempDir(); + DocumentLookupProvider provider = mock(DocumentLookupProvider.class); + when(provider.getDocsAboveSeqNo(anyLong(), any(), any(), any())).thenThrow(new IOException("simulated read failure")); + + String uuid = Translog.createEmptyTranslog(translogPath, SequenceNumbers.NO_OPS_PERFORMED, shardId, primaryTerm.get()); + bootstrapStoreWithMetadata(store, uuid, 3L); + + EngineConfig config = buildDFAEngineConfig(store, translogPath); + EngineCreationFailureException ex = expectThrows( + EngineCreationFailureException.class, + () -> new DataFormatAwareEngine(config, provider) + ); + assertThat(ex.getMessage(), containsString("failed to restore version map")); + assertThat(ex.getCause(), instanceOf(IOException.class)); + } + + /** + * Covers getVersionFromMap unsafe path: versionMap.isUnsafe() triggers refresh + enforceSafeAccess. + */ + @SuppressForbidden(reason = "test needs reflective access to mark versionMap as unsafe") + public void testGetVersionFromMapUnsafePathTriggersRefresh() throws Exception { + DocumentLookupProvider provider = mockLookupProvider(); + try (DataFormatAwareEngine engine = createDFAEngineWithLookupProvider(store, createTempDir(), provider)) { + engine.index(indexOp(createParsedDocWithInput("1", null))); + + java.lang.reflect.Field vmField = DataFormatAwareEngine.class.getDeclaredField("versionMap"); + vmField.setAccessible(true); + LiveVersionMap versionMap = (LiveVersionMap) vmField.get(engine); + java.lang.reflect.Field mapsField = LiveVersionMap.class.getDeclaredField("maps"); + mapsField.setAccessible(true); + Object maps = mapsField.get(versionMap); + java.lang.reflect.Field currentField = maps.getClass().getDeclaredField("current"); + currentField.setAccessible(true); + Object current = currentField.get(maps); + java.lang.reflect.Method markUnsafe = current.getClass().getDeclaredMethod("markAsUnsafe"); + markUnsafe.setAccessible(true); + markUnsafe.invoke(current); + + assertTrue("versionMap should be unsafe", versionMap.isUnsafe()); + + long genBefore; + try (GatedCloseable ref = engine.acquireSnapshot()) { + genBefore = ref.get().getGeneration(); + } + + getByIdLookup(engine, realtimeGet("1")); + + long genAfter; + try (GatedCloseable ref = engine.acquireSnapshot()) { + genAfter = ref.get().getGeneration(); + } + assertThat("refresh should have been triggered by unsafe versionMap", genAfter, greaterThan(genBefore)); + assertFalse("versionMap should be safe after refresh", versionMap.isUnsafe()); + } + } + + /** + * DFAE port of InternalEngineTests.testVersionMapAfterAutoIDDocument (delete step dropped — + * engine.delete() is unsupported in this harness). Verifies LiveVersionMap safe-access + * transitions: optimized append-only skips the map; an explicit-id index enforces safe access + * and stores the entry; safe access is carried over across refresh. + */ + @SuppressForbidden(reason = "test needs reflective access to the engine's versionMap field") + public void testVersionMapAfterAutoIDDocument() throws Exception { + try (DataFormatAwareEngine engine = createDFAEngine(store, createTempDir())) { + java.lang.reflect.Field vmField = DataFormatAwareEngine.class.getDeclaredField("versionMap"); + vmField.setAccessible(true); + LiveVersionMap versionMap = (LiveVersionMap) vmField.get(engine); + + // Optimized append-only: no safe access enforced, nothing stored in the current map. + engine.index(appendOnlyOp(createParsedDocWithInput("1", null))); + assertFalse(versionMap.isSafeAccessRequired()); + assertTrue(versionMap.getAllCurrent().isEmpty()); + + // Explicit-id index of the same id: enforces safe access and stores the entry. + engine.index(indexOp(createParsedDocWithInput("1", null))); + assertTrue(versionMap.isSafeAccessRequired()); + assertEquals(1, versionMap.getAllCurrent().size()); + + // Refresh rotates the map; safe access is carried over. + engine.refresh("test"); + assertTrue(versionMap.isSafeAccessRequired()); + + // Append-only under active safe access still stores the entry. + engine.index(appendOnlyOp(createParsedDocWithInput("2", null))); + assertEquals(1, versionMap.getAllCurrent().size()); + } + } + + /** + * A: write-path version conflict resolved from the versionMap. EXTERNAL version=5 then a stale + * EXTERNAL version=3 hits isVersionConflictForWrites (IndexingStrategyPlanner) → conflict. + */ + public void testIndexVersionConflictFromVersionMap() throws IOException { + try (DataFormatAwareEngine engine = createDFAEngine(store, createTempDir())) { + Engine.IndexResult first = engine.index(indexOpWithVersion(createParsedDocWithInput("1", null), 5L, VersionType.EXTERNAL)); + assertThat(first.getResultType(), equalTo(Engine.Result.Type.SUCCESS)); + Engine.IndexResult stale = engine.index(indexOpWithVersion(createParsedDocWithInput("1", null), 3L, VersionType.EXTERNAL)); + assertThat(stale.getResultType(), equalTo(Engine.Result.Type.FAILURE)); + assertThat(stale.getFailure(), instanceOf(VersionConflictEngineException.class)); + } + } + + /** + * B: optimistic-concurrency conflict on the write path. After indexing id "1" (seqNo 0), an index + * with a non-matching ifSeqNo hits the seqNo/term conflict branch (IndexingStrategyPlanner) → conflict. + */ + public void testIndexIfSeqNoConflict() throws IOException { + try (DataFormatAwareEngine engine = createDFAEngine(store, createTempDir())) { + engine.index(indexOp(createParsedDocWithInput("1", null))); + Engine.IndexResult result = engine.index(indexOpWithIfSeqNo(createParsedDocWithInput("1", null), 99L, primaryTerm.get())); + assertThat(result.getResultType(), equalTo(Engine.Result.Type.FAILURE)); + assertThat(result.getFailure(), instanceOf(VersionConflictEngineException.class)); + } + } + + /** + * C: optimized append-only ops never enter safe-access mode — they leave the current map empty + * and mark it unsafe (LiveVersionMap.maybePutIndexUnderLock else-branch). + */ + @SuppressForbidden(reason = "test needs reflective access to the engine's versionMap field") + public void testAppendOnlyMarksVersionMapUnsafe() throws Exception { + try (DataFormatAwareEngine engine = createDFAEngine(store, createTempDir())) { + java.lang.reflect.Field vmField = DataFormatAwareEngine.class.getDeclaredField("versionMap"); + vmField.setAccessible(true); + LiveVersionMap versionMap = (LiveVersionMap) vmField.get(engine); + + int n = randomIntBetween(2, 5); + for (int i = 0; i < n; i++) { + engine.index(appendOnlyOp(createParsedDocWithInput(Integer.toString(i), null))); + } + assertFalse(versionMap.isSafeAccessRequired()); + assertTrue(versionMap.getAllCurrent().isEmpty()); + assertTrue(versionMap.isUnsafe()); + } + } + + /** + * Covers getById paths: delete check, version conflict, seqNo/primaryTerm conflict. + * Uses reflection to inject a DeleteVersionValue into versionMap since engine.delete() is unsupported. + */ + @SuppressForbidden(reason = "test needs reflective access to inject DeleteVersionValue into versionMap") + public void testGetByIdDeleteAndConflictPaths() throws Exception { + DocumentLookupProvider provider = mockLookupProvider(); + try (DataFormatAwareEngine engine = createDFAEngineWithLookupProvider(store, createTempDir(), provider)) { + // Index a doc so versionMap has an IndexVersionValue + engine.index(indexOp(createParsedDocWithInput("1", null))); + + // --- Path 1: versionValue.isDelete() --- + // Inject a DeleteVersionValue into versionMap via reflection + java.lang.reflect.Field vmField = DataFormatAwareEngine.class.getDeclaredField("versionMap"); + vmField.setAccessible(true); + LiveVersionMap versionMap = (LiveVersionMap) vmField.get(engine); + org.apache.lucene.util.BytesRef uid = new Term(IdFieldMapper.NAME, Uid.encodeId("1")).bytes(); + try (org.opensearch.common.lease.Releasable ignored = versionMap.acquireLock(uid)) { + versionMap.putDeleteUnderLock(uid, new DeleteVersionValue(2L, 1L, 1L, System.currentTimeMillis())); + } + + Engine.Get getForDelete = realtimeGet("1"); + DocumentLookupResult deleteResult = getByIdLookup(engine, getForDelete); + assertFalse("deleted doc should return not found", deleteResult.exists()); + + // --- Path 2: version conflict for reads --- + // Put back an IndexVersionValue so we can test version conflict + try (org.opensearch.common.lease.Releasable ignored = versionMap.acquireLock(uid)) { + versionMap.putIndexUnderLock(uid, new IndexVersionValue(null, 5L, 0L, 1L)); + } + + Engine.Get getWithConflict = realtimeGet("1").version(3L).versionType(VersionType.EXTERNAL); + expectThrows(VersionConflictEngineException.class, () -> getByIdLookup(engine, getWithConflict)); + + // --- Path 3: seqNo/primaryTerm conflict (seqNo matches, primaryTerm doesn't) --- + Engine.Get getWithSeqNoConflict = realtimeGet("1").setIfSeqNo(99L).setIfPrimaryTerm(1L); + expectThrows(VersionConflictEngineException.class, () -> getByIdLookup(engine, getWithSeqNoConflict)); + + // --- Path 4: primaryTerm mismatch (seqNo matches but primaryTerm doesn't) --- + // versionMap has seqNo=0, primaryTerm=1 from the IndexVersionValue above + try (org.opensearch.common.lease.Releasable ignored = versionMap.acquireLock(uid)) { + versionMap.putIndexUnderLock(uid, new IndexVersionValue(null, 5L, 0L, 1L)); + } + Engine.Get getWithPtConflict = realtimeGet("1").setIfSeqNo(0L).setIfPrimaryTerm(999L); + expectThrows(VersionConflictEngineException.class, () -> getByIdLookup(engine, getWithPtConflict)); + + // --- Path 5: seqNo and primaryTerm both match — no conflict, falls through --- + Engine.Get getNoConflict = realtimeGet("1").setIfSeqNo(0L).setIfPrimaryTerm(1L); + DocumentLookupResult noConflictResult = getByIdLookup(engine, getNoConflict); + assertNotNull("should not throw when seqNo and primaryTerm match", noConflictResult); + + // --- Path 6: GC deletes in resolveDocVersion --- + // Inject an expired DeleteVersionValue (time far in the past) so gc_deletes nullifies it + try (org.opensearch.common.lease.Releasable ignored = versionMap.acquireLock(uid)) { + versionMap.putDeleteUnderLock(uid, new DeleteVersionValue(3L, 2L, 1L, 1L)); // time=1ms (epoch start, expired) + } + // Index same doc again — resolveDocVersion finds expired DeleteVersionValue, nullifies it, + // treats as new doc (version NOT_FOUND) + Engine.IndexResult gcResult = engine.index(indexOp(createParsedDocWithInput("1", null))); + assertThat(gcResult.getResultType(), equalTo(Engine.Result.Type.SUCCESS)); + + // --- Path 7: translog readOperation returns null (bogus location) --- + // Covers the negative path: if (operation != null) is false, falls through to refreshIfNeeded + Translog.Location bogusLocation = new Translog.Location(999L, 0L, 1); + try (org.opensearch.common.lease.Releasable ignored = versionMap.acquireLock(uid)) { + versionMap.putIndexUnderLock(uid, new IndexVersionValue(bogusLocation, 6L, 3L, 1L)); + } + Engine.Get getWithBadTranslog = realtimeGet("1"); + DocumentLookupResult translogNullResult = getByIdLookup(engine, getWithBadTranslog); + assertNotNull("should fall through when translog returns null", translogNullResult); + } + } } diff --git a/server/src/test/java/org/opensearch/index/engine/DataFormatAwareReadOnlyEngineTests.java b/server/src/test/java/org/opensearch/index/engine/DataFormatAwareReadOnlyEngineTests.java index 9ebba95a6b7e8..a6a0e709f217d 100644 --- a/server/src/test/java/org/opensearch/index/engine/DataFormatAwareReadOnlyEngineTests.java +++ b/server/src/test/java/org/opensearch/index/engine/DataFormatAwareReadOnlyEngineTests.java @@ -36,6 +36,7 @@ import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import org.opensearch.index.engine.exec.coord.CatalogSnapshotManager; import org.opensearch.index.engine.exec.coord.DataformatAwareCatalogSnapshot; +import org.opensearch.index.get.DocumentLookupResult; import org.opensearch.index.seqno.RetentionLeases; import org.opensearch.index.seqno.SequenceNumbers; import org.opensearch.index.shard.ShardPath; @@ -43,6 +44,7 @@ import org.opensearch.index.store.Store; import org.opensearch.index.translog.Translog; import org.opensearch.index.translog.TranslogConfig; +import org.opensearch.plugins.DocumentLookupProvider; import org.opensearch.plugins.PluginsService; import org.opensearch.plugins.SearchBackEndPlugin; import org.opensearch.test.DummyShardLock; @@ -69,6 +71,7 @@ import static org.opensearch.index.engine.EngineTestCase.tombstoneDocSupplier; import static org.hamcrest.Matchers.instanceOf; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -844,4 +847,138 @@ public void testNoOpDeletionPolicyNeverDeletesSnapshots() { assertTrue("onInit with empty list should return empty", policy.onInit(List.of()).isEmpty()); assertTrue("onCommit with empty list should return empty", policy.onCommit(List.of()).isEmpty()); } + + // ---------- getById Tests ---------- + + private Engine.Get realtimeGet(String id) { + return new Engine.Get(true, true, id, new org.apache.lucene.index.Term("_id", org.opensearch.index.mapper.Uid.encodeId(id))); + } + + private DataFormatAwareReadOnlyEngine createReadOnlyEngineWithProvider(org.opensearch.plugins.DocumentLookupProvider provider) + throws IOException { + String translogUUID = UUID.randomUUID().toString(); + String historyUUID = UUID.randomUUID().toString(); + bootstrapStoreWithMetadata(store, translogUUID, historyUUID); + Path translogPath = createTempDir().resolve("translog"); + TranslogConfig translogConfig = new TranslogConfig( + shardId, + translogPath, + warmIndexSettings(), + BigArrays.NON_RECYCLING_INSTANCE, + "", + false + ); + DataFormatRegistry registry = createMockRegistry(); + CommitterFactory committerFactory = config -> new InMemoryCommitter(store) { + @Override + public List listCommittedSnapshots() { + Map userData = new HashMap<>(); + userData.put(SequenceNumbers.LOCAL_CHECKPOINT_KEY, Long.toString(SequenceNumbers.NO_OPS_PERFORMED)); + userData.put(SequenceNumbers.MAX_SEQ_NO, Long.toString(SequenceNumbers.NO_OPS_PERFORMED)); + userData.put(Translog.TRANSLOG_UUID_KEY, UUID.randomUUID().toString()); + userData.put(Engine.HISTORY_UUID_KEY, UUID.randomUUID().toString()); + DataformatAwareCatalogSnapshot snapshot = (DataformatAwareCatalogSnapshot) CatalogSnapshotManager.createInitialSnapshot( + 0L, + 0L, + 0L, + List.of(new Segment(1L, Map.of())), + -1L, + userData + ); + snapshot.setLastCommitInfo("segments_1", 1L, 0L); + return List.of(snapshot); + } + }; + EngineConfig config = new EngineConfig.Builder().shardId(shardId) + .threadPool(threadPool) + .indexSettings(warmIndexSettings()) + .store(store) + .mergePolicy(org.apache.lucene.index.NoMergePolicy.INSTANCE) + .translogConfig(translogConfig) + .flushMergesAfter(TimeValue.timeValueMinutes(5)) + .externalRefreshListener(List.of()) + .internalRefreshListener(List.of()) + .globalCheckpointSupplier(() -> SequenceNumbers.NO_OPS_PERFORMED) + .retentionLeasesSupplier(() -> RetentionLeases.EMPTY) + .primaryTermSupplier(primaryTerm::get) + .tombstoneDocSupplier(tombstoneDocSupplier()) + .dataFormatRegistry(registry) + .committerFactory(committerFactory) + .readOnlyReplica(false) + .build(); + return new DataFormatAwareReadOnlyEngine(config, provider); + } + + /** Calls the new getById(Engine.Get, searcherFactory) and unwraps the pre-materialized lookup. + * DFA engines ignore the searcher factory, so a no-op factory is supplied. */ + private static DocumentLookupResult getByIdLookup(DataFormatAwareReadOnlyEngine engine, Engine.Get get) throws IOException { + Engine.GetResult result = engine.getById(get, (source, scope) -> null); + return result.exists() ? ((DocumentLookupResult.PreMaterialized) result).lookup() : DocumentLookupResult.notFound(get.id()); + } + + public void testGetByIdThrowsWhenNoProvider() throws IOException { + try (DataFormatAwareReadOnlyEngine engine = createReadOnlyEngine()) { + Engine.Get get = realtimeGet("1"); + expectThrows(UnsupportedOperationException.class, () -> getByIdLookup(engine, get)); + } + } + + public void testGetByIdReturnsNotFoundOnEmptyCatalog() throws IOException { + DocumentLookupProvider provider = mock(DocumentLookupProvider.class); + String translogUUID = UUID.randomUUID().toString(); + String historyUUID = UUID.randomUUID().toString(); + bootstrapStoreWithMetadata(store, translogUUID, historyUUID); + EngineConfig config = buildConfig(store, warmIndexSettings(), null); + try (DataFormatAwareReadOnlyEngine engine = new DataFormatAwareReadOnlyEngine(config, provider)) { + Engine.Get get = realtimeGet("1"); + DocumentLookupResult result = getByIdLookup(engine, get); + assertFalse("should return not found when catalog has no segments", result.exists()); + } + } + + public void testGetByIdReturnsDocWhenNoConflict() throws IOException { + DocumentLookupProvider provider = mock(DocumentLookupProvider.class); + when(provider.getById(any(), any(), any(), any())).thenReturn( + new DocumentLookupResult("1", 1L, true, null, 0L, 1L, Map.of(), Map.of()) + ); + try (DataFormatAwareReadOnlyEngine engine = createReadOnlyEngineWithProvider(provider)) { + Engine.Get get = realtimeGet("1"); + DocumentLookupResult result = getByIdLookup(engine, get); + assertTrue("doc should be found", result.exists()); + assertEquals(1L, result.version()); + } + } + + public void testGetByIdThrowsVersionConflictForReads() throws IOException { + DocumentLookupProvider provider = mock(DocumentLookupProvider.class); + when(provider.getById(any(), any(), any(), any())).thenReturn( + new DocumentLookupResult("1", 5L, true, null, 0L, 1L, Map.of(), Map.of()) + ); + try (DataFormatAwareReadOnlyEngine engine = createReadOnlyEngineWithProvider(provider)) { + Engine.Get get = realtimeGet("1").version(3L).versionType(org.opensearch.index.VersionType.EXTERNAL); + expectThrows(VersionConflictEngineException.class, () -> getByIdLookup(engine, get)); + } + } + + public void testGetByIdThrowsOnSeqNoMismatch() throws IOException { + DocumentLookupProvider provider = mock(DocumentLookupProvider.class); + when(provider.getById(any(), any(), any(), any())).thenReturn( + new DocumentLookupResult("1", 1L, true, null, 5L, 1L, Map.of(), Map.of()) + ); + try (DataFormatAwareReadOnlyEngine engine = createReadOnlyEngineWithProvider(provider)) { + Engine.Get get = realtimeGet("1").setIfSeqNo(99L).setIfPrimaryTerm(1L); + expectThrows(VersionConflictEngineException.class, () -> getByIdLookup(engine, get)); + } + } + + public void testGetByIdThrowsOnPrimaryTermMismatch() throws IOException { + DocumentLookupProvider provider = mock(DocumentLookupProvider.class); + when(provider.getById(any(), any(), any(), any())).thenReturn( + new DocumentLookupResult("1", 1L, true, null, 5L, 1L, Map.of(), Map.of()) + ); + try (DataFormatAwareReadOnlyEngine engine = createReadOnlyEngineWithProvider(provider)) { + Engine.Get get = realtimeGet("1").setIfSeqNo(5L).setIfPrimaryTerm(99L); + expectThrows(VersionConflictEngineException.class, () -> getByIdLookup(engine, get)); + } + } } diff --git a/server/src/test/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshotTests.java b/server/src/test/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshotTests.java index 8a1c94256ebea..4b3d02da99604 100644 --- a/server/src/test/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshotTests.java +++ b/server/src/test/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshotTests.java @@ -406,6 +406,47 @@ public void testClonePreservesUserData() { assertEquals(5L, cloned.getVersion()); } + // --- findFileSet --- + + public void testFindFileSet_matchesGeneration() { + WriterFileSet pset = new WriterFileSet("/tmp/pq", 5L, Set.of("f.parquet"), 100, 0L); + DataformatAwareCatalogSnapshot snapshot = new DataformatAwareCatalogSnapshot( + 1L, + 1L, + 1L, + List.of(new Segment(5L, Map.of("parquet", pset))), + 0L, + Map.of() + ); + assertEquals(pset, snapshot.findFileSet("parquet", 5L)); + } + + public void testFindFileSet_noMatchingGeneration_returnsNull() { + WriterFileSet pset = new WriterFileSet("/tmp/pq", 5L, Set.of("f.parquet"), 100, 0L); + DataformatAwareCatalogSnapshot snapshot = new DataformatAwareCatalogSnapshot( + 1L, + 1L, + 1L, + List.of(new Segment(5L, Map.of("parquet", pset))), + 0L, + Map.of() + ); + assertNull(snapshot.findFileSet("parquet", 99L)); + } + + public void testFindFileSet_skipsEmptyFiles_returnsNull() { + WriterFileSet empty = new WriterFileSet("/tmp/pq", 5L, Set.of(), 0, 0L); + DataformatAwareCatalogSnapshot snapshot = new DataformatAwareCatalogSnapshot( + 1L, + 1L, + 1L, + List.of(new Segment(5L, Map.of("parquet", empty))), + 0L, + Map.of() + ); + assertNull(snapshot.findFileSet("parquet", 5L)); + } + // --- helpers --- private WriterFileSet randomWriterFileSet(String format) { diff --git a/server/src/test/java/org/opensearch/index/get/DocumentLookupResultTests.java b/server/src/test/java/org/opensearch/index/get/DocumentLookupResultTests.java new file mode 100644 index 0000000000000..7a77e25959cef --- /dev/null +++ b/server/src/test/java/org/opensearch/index/get/DocumentLookupResultTests.java @@ -0,0 +1,68 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.index.get; + +import org.opensearch.common.lucene.uid.Versions; +import org.opensearch.core.common.bytes.BytesArray; +import org.opensearch.core.common.bytes.BytesReference; +import org.opensearch.index.engine.Engine; +import org.opensearch.index.seqno.SequenceNumbers; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Map; + +public class DocumentLookupResultTests extends OpenSearchTestCase { + + public void testNotFound() { + DocumentLookupResult r = DocumentLookupResult.notFound("x"); + assertFalse(r.exists()); + assertEquals("x", r.id()); + assertEquals(Versions.NOT_FOUND, r.version()); + assertEquals(SequenceNumbers.UNASSIGNED_SEQ_NO, r.seqNo()); + assertEquals(SequenceNumbers.UNASSIGNED_PRIMARY_TERM, r.primaryTerm()); + assertNull(r.source()); + assertTrue(r.documentFields().isEmpty()); + assertTrue(r.metadataFields().isEmpty()); + } + + public void testNullFieldMapsNormalizedToEmpty() { + DocumentLookupResult r = new DocumentLookupResult("1", 1L, true, null, 0L, 1L, null, null); + assertNotNull(r.documentFields()); + assertNotNull(r.metadataFields()); + assertTrue(r.documentFields().isEmpty()); + assertTrue(r.metadataFields().isEmpty()); + } + + public void testEqualsAndHashCode() { + BytesReference src = new BytesArray("{\"f\":1}"); + DocumentLookupResult a = new DocumentLookupResult("1", 2L, true, src, 3L, 1L, Map.of(), Map.of()); + DocumentLookupResult b = new DocumentLookupResult("1", 2L, true, new BytesArray("{\"f\":1}"), 3L, 1L, Map.of(), Map.of()); + assertEquals(a, b); + assertEquals(a.hashCode(), b.hashCode()); + + DocumentLookupResult differentVersion = new DocumentLookupResult("1", 9L, true, src, 3L, 1L, Map.of(), Map.of()); + assertNotEquals(a, differentVersion); + } + + public void testPreMaterializedShape() { + DocumentLookupResult r = new DocumentLookupResult("1", 5L, true, null, 7L, 2L, Map.of(), Map.of()); + Engine.GetResult gr = r.toGetResult(); + + assertTrue(gr instanceof DocumentLookupResult.PreMaterialized); + assertTrue(gr.exists()); + assertEquals(5L, gr.version()); + assertSame(r, ((DocumentLookupResult.PreMaterialized) gr).lookup()); + + // PreMaterialized has no Lucene searcher/docId — these must fail fast, not be dereferenced. + expectThrows(UnsupportedOperationException.class, gr::searcher); + expectThrows(UnsupportedOperationException.class, gr::docIdAndVersion); + + gr.close(); // no-op, must not throw + } +} diff --git a/server/src/test/java/org/opensearch/index/shard/ShardGetServiceTests.java b/server/src/test/java/org/opensearch/index/shard/ShardGetServiceTests.java index e261a244742cc..97fa0cfc7253b 100644 --- a/server/src/test/java/org/opensearch/index/shard/ShardGetServiceTests.java +++ b/server/src/test/java/org/opensearch/index/shard/ShardGetServiceTests.java @@ -33,19 +33,28 @@ import org.opensearch.Version; import org.opensearch.cluster.metadata.IndexMetadata; +import org.opensearch.common.SuppressForbidden; +import org.opensearch.common.document.DocumentField; import org.opensearch.common.settings.Settings; +import org.opensearch.core.common.bytes.BytesArray; +import org.opensearch.core.common.bytes.BytesReference; import org.opensearch.core.index.shard.ShardId; import org.opensearch.core.xcontent.MediaTypeRegistry; import org.opensearch.index.VersionType; import org.opensearch.index.engine.Engine; import org.opensearch.index.engine.VersionConflictEngineException; +import org.opensearch.index.get.DocumentLookupResult; import org.opensearch.index.get.GetResult; +import org.opensearch.index.get.ShardGetService; import org.opensearch.index.mapper.MapperService; import org.opensearch.index.mapper.RoutingFieldMapper; import org.opensearch.search.fetch.subphase.FetchSourceContext; import java.io.IOException; +import java.lang.reflect.Method; import java.nio.charset.StandardCharsets; +import java.util.Collections; +import java.util.Map; import static org.opensearch.index.seqno.SequenceNumbers.UNASSIGNED_PRIMARY_TERM; import static org.opensearch.index.seqno.SequenceNumbers.UNASSIGNED_SEQ_NO; @@ -217,4 +226,55 @@ public void testTypelessGetForUpdate() throws IOException { closeShards(shard); } + + @SuppressForbidden(reason = "reflective access to test private buildFromLookup method") + public void testBuildFromLookupSourceFiltering() throws Exception { + Settings settings = Settings.builder() + .put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT) + .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1) + .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1) + .build(); + IndexMetadata metadata = IndexMetadata.builder("test") + .putMapping("{ \"properties\": { \"foo\": { \"type\": \"text\"}}}") + .settings(settings) + .primaryTerm(0, 1) + .build(); + IndexShard primary = newShard(new ShardId(metadata.getIndex(), 0), true, "n1", metadata, null); + recoverShardFromStore(primary); + + ShardGetService getService = primary.getService(); + Method buildFromLookup = ShardGetService.class.getDeclaredMethod( + "buildFromLookup", + DocumentLookupResult.class, + FetchSourceContext.class + ); + buildFromLookup.setAccessible(true); + + BytesReference source = new BytesArray("{\"foo\":\"bar\",\"baz\":\"qux\"}"); + DocumentField docField = new DocumentField("foo", Collections.singletonList("bar")); + DocumentLookupResult lookup = new DocumentLookupResult("1", 3L, true, source, 5L, 1L, Map.of("foo", docField), Map.of()); + + // Case 1: FETCH_SOURCE — source passes through unfiltered + GetResult result = (GetResult) buildFromLookup.invoke(getService, lookup, FetchSourceContext.FETCH_SOURCE); + assertTrue(result.isExists()); + assertEquals("1", result.getId()); + assertEquals(3L, result.getVersion()); + assertEquals(5L, result.getSeqNo()); + assertEquals(1L, result.getPrimaryTerm()); + assertEquals("{\"foo\":\"bar\",\"baz\":\"qux\"}", new String(result.source(), StandardCharsets.UTF_8)); + assertEquals("bar", result.getFields().get("foo").getValue()); + + // Case 2: DO_NOT_FETCH_SOURCE — source is null + GetResult noSourceResult = (GetResult) buildFromLookup.invoke(getService, lookup, FetchSourceContext.DO_NOT_FETCH_SOURCE); + assertNull(noSourceResult.source()); + assertTrue(noSourceResult.isExists()); + assertEquals("1", noSourceResult.getId()); + + // Case 3: includes filter — only matching fields in source + FetchSourceContext includesOnly = new FetchSourceContext(true, new String[] { "foo" }, new String[0]); + GetResult filteredResult = (GetResult) buildFromLookup.invoke(getService, lookup, includesOnly); + assertEquals("{\"foo\":\"bar\"}", new String(filteredResult.source(), StandardCharsets.UTF_8)); + + closeShards(primary); + } } diff --git a/server/src/test/java/org/opensearch/plugins/DocumentLookupProviderTests.java b/server/src/test/java/org/opensearch/plugins/DocumentLookupProviderTests.java new file mode 100644 index 0000000000000..593f60e6201cf --- /dev/null +++ b/server/src/test/java/org/opensearch/plugins/DocumentLookupProviderTests.java @@ -0,0 +1,34 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugins; + +import org.opensearch.core.index.Index; +import org.opensearch.index.engine.exec.DocumentMetadataResolver; +import org.opensearch.index.get.DocumentLookupResult; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; + +public class DocumentLookupProviderTests extends OpenSearchTestCase { + + private static final Index INDEX = new Index("idx", "uuid"); + + /** A provider that overrides only the abstract getById; getVersionMetadata/getDocsAboveSeqNo use the SPI defaults. */ + private static final DocumentLookupProvider DEFAULTS = (get, reader, index, resolver) -> DocumentLookupResult.notFound("x"); + + public void testGetVersionMetadataDefaultsToNotFound() throws IOException { + DocumentLookupResult result = DEFAULTS.getVersionMetadata("id", null, INDEX, DocumentMetadataResolver.NOOP); + assertFalse(result.exists()); + assertEquals("id", result.id()); + } + + public void testGetDocsAboveSeqNoDefaultsToEmpty() throws IOException { + assertTrue(DEFAULTS.getDocsAboveSeqNo(0L, null, INDEX, DocumentMetadataResolver.NOOP).isEmpty()); + } +} From 8d707367f9052a6350539d5255b70c242689f1a5 Mon Sep 17 00:00:00 2001 From: Arpit Bandejiya Date: Fri, 19 Jun 2026 18:11:59 +0530 Subject: [PATCH 2/2] Add internalSearch flag for native get-by-id (no Substrait) Route the pluggable-dataformat get-by-id and seq-no-scan through the existing df_execute_query FFM call with an internalSearch flag instead of synthesizing SQL->Substrait in GetService. When the flag is set, the native side builds a DataFrame plan with a pushed-down filter (__row_id__ = n for get-by-id, _seq_no > floor for the recovery scan), executes it, and returns a stream drained through the same df_stream_next/df_stream_close path as a normal query. - ffm.rs: df_execute_query gains internal_search_mode + internal_search_bound - datafusion_query_config.rs: InternalSearch enum + from_wire + unit test - api.rs: execute_internal_search (DataFrame .filter, single-file assert, pushdown on); early dispatch from execute_query - helpers.rs (new): consolidate shared query setup used by both the Substrait and internal-search paths -- build_query_runtime_env (list-files cache + per-query memory pool overlay + object-store registration), build_session_context (with register_udfs flag), new_shard_query_context, register_listing_table, execute_plan_to_stream_ptr, run_cancellable_to_handle_ptr, wrap_stream_as_handle; repoint query_executor/indexed_executor/session_context - NativeBridge.java: executeQueryAsync gains internalSearchMode/Bound + INTERNAL_SEARCH_* constants; FFM descriptor updated - GetService.java: drop sqlToSubstrait + WireConfigSnapshot; call the flagged path --- .../rust/src/api.rs | 4 + .../rust/src/datafusion_query_config.rs | 47 ++++++++ .../rust/src/ffm.rs | 9 ++ .../rust/src/query_executor.rs | 70 ++++++++++-- .../opensearch/be/datafusion/GetService.java | 104 ++++++++---------- .../be/datafusion/nativelib/NativeBridge.java | 55 +++++++-- 6 files changed, 216 insertions(+), 73 deletions(-) diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs index 28b2717e7de76..3052dab14caa8 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs @@ -841,6 +841,7 @@ pub async unsafe fn execute_query( manager: &RuntimeManager, context_id: i64, query_config: crate::datafusion_query_config::DatafusionQueryConfig, + internal_search: crate::datafusion_query_config::InternalSearch, ) -> Result { let shard_view = &*(shard_view_ptr as *const ShardView); let runtime = &*(runtime_ptr as *const DataFusionRuntime); @@ -911,6 +912,8 @@ pub async unsafe fn execute_query( // - IndexedPredicateOnly → indexed path (position-based row IDs) // - None → vanilla path (no row ID computation) // 3. Neither → vanilla path + // Engine-internal point lookups pass an empty plan, so is_indexed/has_row_id are both + // false and this naturally resolves to the vanilla ListingTable path. let use_indexed = is_indexed || (has_row_id && effective_config.query_strategy != crate::datafusion_query_config::QueryStrategy::ListingTable); @@ -941,6 +944,7 @@ pub async unsafe fn execute_query( phantom_corrector, &shard_view.sort_fields, &shard_view.sort_orders, + internal_search, ).await } }; diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/datafusion_query_config.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/datafusion_query_config.rs index 5d8cdb1df44dd..0f819690ff380 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/datafusion_query_config.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/datafusion_query_config.rs @@ -28,6 +28,41 @@ pub enum QueryStrategy { IndexedPredicateOnly, } +/// Engine-internal point lookup driven through the normal `df_execute_query` +/// entry point. When active, the Substrait `plan_ptr` is ignored and the plan +/// is built natively via the DataFrame API with a single pushed-down filter on +/// a stored reserved column — no Substrait, no planner round-trip. Used by the +/// pluggable-dataformat get-by-id path (`GetService`), not by user search. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum InternalSearch { + /// Not an internal lookup — decode `plan_ptr` as Substrait as usual. + Off, + /// Get-by-row-id: `__row_id__ = bound`, single row. `bound` is the physical + /// row position resolved from the secondary (Lucene) index. + ByRowId(i64), + /// Seq-no scan: `_seq_no > bound`, projecting only id/seq/term/version. + /// Used by version-map restore on crash recovery. + SeqNoAbove(i64), +} + +impl InternalSearch { + /// Decodes the FFM wire pair `(mode, bound)`. `mode`: 0 = Off, 1 = ByRowId, + /// 2 = SeqNoAbove. Any other value is treated as Off (forward-compatible). + pub fn from_wire(mode: i64, bound: i64) -> Self { + match mode { + 1 => InternalSearch::ByRowId(bound), + 2 => InternalSearch::SeqNoAbove(bound), + _ => InternalSearch::Off, + } + } + + /// Whether this is an engine-internal point lookup (i.e. not [`InternalSearch::Off`], + /// the normal user-search path). + pub fn is_internal_search(self) -> bool { + !matches!(self, InternalSearch::Off) + } +} + /// Query-scoped configuration. Owned by value after FFM decode. #[derive(Debug, Clone)] pub struct DatafusionQueryConfig { @@ -316,6 +351,18 @@ mod tests { unsafe { DatafusionQueryConfig::from_ffm_ptr(0) }; } + #[test] + fn internal_search_from_wire_decodes_modes() { + assert_eq!(InternalSearch::from_wire(0, 99), InternalSearch::Off); + assert_eq!(InternalSearch::from_wire(1, 42), InternalSearch::ByRowId(42)); + assert_eq!(InternalSearch::from_wire(2, 7), InternalSearch::SeqNoAbove(7)); + // Unknown modes are forward-compatible: treated as Off, bound ignored. + assert_eq!(InternalSearch::from_wire(3, 5), InternalSearch::Off); + assert!(!InternalSearch::Off.is_internal_search()); + assert!(InternalSearch::ByRowId(0).is_internal_search()); + assert!(InternalSearch::SeqNoAbove(0).is_internal_search()); + } + #[test] fn wire_decode_round_trips_all_fields() { let wire = WireDatafusionQueryConfig { diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs index 55e8719813e64..cf98dc504e66d 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs @@ -46,6 +46,7 @@ fn timed_block_on( use crate::api; use crate::api::DataFusionRuntime; use crate::cache; +use crate::datafusion_query_config::InternalSearch; use crate::custom_cache_manager::CustomCacheManager; use crate::eviction_policy::PolicyType; use crate::runtime_manager::RuntimeManager; @@ -312,6 +313,12 @@ pub unsafe extern "C" fn df_execute_query( context_id: i64, // Pointer to a `WireDatafusionQueryConfig` query_config_ptr: i64, + // Engine-internal point lookup. 0 = normal query (decode `plan_ptr` as Substrait); + // 1 = get-by-row-id (`__row_id__ = internal_search_bound`); 2 = seq-no scan + // (`_seq_no > internal_search_bound`). When non-zero, `plan_ptr`/`plan_len` are + // ignored and the plan is built natively via the DataFrame API — no Substrait. + internal_search_mode: i64, + internal_search_bound: i64, ) -> i64 { let mgr = get_rt_manager()?; let table_name = str_from_raw(table_name_ptr, table_name_len) @@ -319,6 +326,7 @@ pub unsafe extern "C" fn df_execute_query( let plan_bytes = slice::from_raw_parts(plan_ptr, plan_len as usize); let query_config = crate::datafusion_query_config::DatafusionQueryConfig::from_ffm_ptr(query_config_ptr); + let internal_search = InternalSearch::from_wire(internal_search_mode, internal_search_bound); // Copy the plan bytes so the spawned future can own them (`cpu_executor.spawn` // requires `'static`). The `shard_view_ptr`, `runtime_ptr` are raw pointers // held live by the caller for the duration of the FFM downcall — safe to @@ -346,6 +354,7 @@ pub unsafe extern "C" fn df_execute_query( &mgr_for_inner, context_id, query_config, + internal_search, ) .await }); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs index ba2347b70f093..6f1e05018ebf1 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs @@ -55,6 +55,7 @@ pub async fn execute_query( phantom_corrector: Option>, sort_fields: &[String], sort_orders: &[String], + internal_search: crate::datafusion_query_config::InternalSearch, ) -> Result { // Build per-query RuntimeEnv with list-files cache pre-populated. let runtime_env = build_query_runtime_env(runtime, &table_path, object_metas.as_ref())?; @@ -161,13 +162,9 @@ pub async fn execute_query( } } - // Decode substrait → logical plan → physical plan → stream - let substrait_plan = Plan::decode(plan_bytes.as_slice()).map_err(|e| { - DataFusionError::Execution(format!("Failed to decode Substrait: {}", e)) - })?; - - let logical_plan = from_substrait_plan(&ctx.state(), &substrait_plan).await?; - let dataframe = ctx.execute_logical_plan(logical_plan).await?; + // Planning: build the query DataFrame (Substrait decode for normal search, native filter for an + // engine-internal point lookup). Physical planning + execution below is shared by both. + let dataframe = build_dataframe(&ctx, &table_name, &plan_bytes, internal_search).await?; let physical_plan = dataframe.create_physical_plan().await?; // Retag any physical-plan output columns whose type tags differ from what Substrait @@ -218,6 +215,65 @@ pub async fn execute_query( Ok(Box::into_raw(Box::new(wrapped)) as i64) } +/// Build the query DataFrame against the table already registered in `ctx`. +/// +/// An engine-internal point lookup (get-by-id / seq-no scan) returns early with a small native +/// filter DataFrame; otherwise this is the standard user-search flow: decode the Substrait plan +/// into a logical plan and execute it. +/// +/// Internal-search filters: +/// - [`InternalSearch::ByRowId`]: `SELECT * WHERE __row_id__ = n LIMIT 1`. `__row_id__` is the +/// physical row position the writer stamps at flush, so equality is order-independent and prunes +/// row-groups/pages via min/max stats. +/// - [`InternalSearch::SeqNoAbove`]: `SELECT _id,_seq_no,_primary_term,_version WHERE _seq_no > f` +/// (version-map restore on recovery; metadata columns only). +async fn build_dataframe( + ctx: &SessionContext, + table_name: &str, + plan_bytes: &[u8], + internal_search: crate::datafusion_query_config::InternalSearch, +) -> Result { + // Engine-internal point lookup — build a native filter DataFrame and return early. + if internal_search.is_internal_search() { + return internal_search_dataframe(ctx, table_name, internal_search).await; + } + + // Standard user-search flow: Substrait → logical plan → DataFrame. + let substrait_plan = Plan::decode(plan_bytes).map_err(|e| { + DataFusionError::Execution(format!("Failed to decode Substrait: {}", e)) + })?; + let logical_plan = from_substrait_plan(&ctx.state(), &substrait_plan).await?; + ctx.execute_logical_plan(logical_plan).await +} + +/// Build the native filter DataFrame for an engine-internal point lookup against the table already +/// registered in `ctx`. Not user search — those go through the Substrait flow in [`build_dataframe`]. +/// +/// - [`InternalSearch::ByRowId`]: `SELECT * WHERE __row_id__ = n LIMIT 1`. `__row_id__` is the +/// physical row position the writer stamps at flush, so equality is order-independent and prunes +/// row-groups/pages via min/max stats. +/// - [`InternalSearch::SeqNoAbove`]: `SELECT _id,_seq_no,_primary_term,_version WHERE _seq_no > f` +/// (version-map restore on recovery; metadata columns only). +/// +/// Panics on [`InternalSearch::Off`] — callers gate on `is_internal_search()` first. +async fn internal_search_dataframe( + ctx: &SessionContext, + table_name: &str, + internal_search: crate::datafusion_query_config::InternalSearch, +) -> Result { + use crate::datafusion_query_config::InternalSearch; + let df = ctx.table(table_name).await?; + match internal_search { + InternalSearch::ByRowId(row_id) => df + .filter(col(crate::ROW_ID_COLUMN_NAME).eq(lit(row_id)))? + .limit(0, Some(1)), + InternalSearch::SeqNoAbove(seq_no_floor) => df + .filter(col("_seq_no").gt(lit(seq_no_floor)))? + .select_columns(&["_id", "_seq_no", "_primary_term", "_version"]), + InternalSearch::Off => unreachable!("internal_search_dataframe called with Off"), + } +} + /// Executes a Substrait plan against a pre-configured SessionContext. /// /// Takes ownership of the handle by value. The ownership transfer (consuming the diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/GetService.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/GetService.java index 2757d3d58c4df..7eeb5165b60d0 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/GetService.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/GetService.java @@ -31,8 +31,6 @@ import java.io.Closeable; import java.io.IOException; -import java.lang.foreign.Arena; -import java.lang.foreign.MemorySegment; import java.util.ArrayList; import java.util.List; import java.util.Map; @@ -81,6 +79,8 @@ static final class NativeBridgeExecutor implements DocumentRowReader, Closeable private static final String GET_BY_ID_TABLE_ALIAS = "_t"; private static final String PARQUET_FORMAT = "parquet"; + /** Empty Substrait plan — the internal-search path builds its plan natively and ignores it. */ + private static final byte[] EMPTY_PLAN = new byte[0]; private final DataFusionPlugin dfPlugin; private final BufferAllocator sharedAllocator = new RootAllocator(64 * 1024 * 1024); @@ -112,18 +112,18 @@ public Map executeSingleRow(long rowId, WriterFileSet parquetSet MonoFileWriterSet segment = MonoFileWriterSet.of(parquetDir, parquetSet.writerGeneration(), parquetFile, 0L); try (ReaderHandle readerHandle = new ReaderHandle(parquetDir, List.of(segment), null, List.of(), List.of())) { long readerPtr = readerHandle.getPointer(); - // TODO: only the OFFSET (rowId) varies between calls — cache the Substrait plan - // per reader and rebind the offset to avoid re-planning on every get-by-id. - String sql = "SELECT * FROM \"" + GET_BY_ID_TABLE_ALIAS + "\" LIMIT 1 OFFSET " + Long.toUnsignedString(rowId); - byte[] substraitPlan = NativeBridge.sqlToSubstrait(readerPtr, GET_BY_ID_TABLE_ALIAS, sql, runtimePtr); - WireConfigSnapshot configSnapshot = WireConfigSnapshot.builder(dfPlugin.getDatafusionSettings().getSnapshot()) - .queryStrategy(1) // ListingTable — bypass indexed executor routing - .build(); - long streamPtr = executeNativeQuery( + // Internal-search get-by-row-id: the native side ignores Substrait and builds a + // DataFrame plan filtering `__row_id__ = rowId` with pushdown enabled. __row_id__ is + // the physical row position the parquet writer stamps at flush (sequential 0..N after + // any index sort) and remaps into the Lucene secondary index, so the resolver's rowId + // equals the row's __row_id__. An equality predicate returns exactly that row + // independent of scan order and lets DataFusion prune row-groups/pages via the + // column's min/max statistics. + long streamPtr = executeInternalSearch( readerPtr, - substraitPlan, runtimePtr, - configSnapshot, + NativeBridge.INTERNAL_SEARCH_BY_ROW_ID, + rowId, "DataFusion get-by-id query failed" ); return readSingleRow(streamPtr); @@ -140,21 +140,14 @@ public List> executeRowsAboveSeqNo(List fileS MonoFileWriterSet writerSet = MonoFileWriterSet.of(parquetDir, parquetSet.writerGeneration(), parquetFile, 0L); try (ReaderHandle readerHandle = new ReaderHandle(parquetDir, List.of(writerSet), null, List.of(), List.of())) { long readerPtr = readerHandle.getPointer(); - // TODO: only the _seq_no floor varies between calls -- cache the Substrait plan - // per reader and rebind the bound to avoid re-planning on every restore scan. - String sql = "SELECT \"_id\", \"_seq_no\", \"_primary_term\", \"_version\" FROM \"" - + GET_BY_ID_TABLE_ALIAS - + "\" WHERE \"_seq_no\" > " - + seqNoFloor; - byte[] substraitPlan = NativeBridge.sqlToSubstrait(readerPtr, GET_BY_ID_TABLE_ALIAS, sql, runtimePtr); - WireConfigSnapshot configSnapshot = WireConfigSnapshot.builder(dfPlugin.getDatafusionSettings().getSnapshot()) - .queryStrategy(1) // ListingTable — bypass indexed executor routing - .build(); - long streamPtr = executeNativeQuery( + // Internal-search seq-no scan: the native side ignores Substrait and builds a + // DataFrame plan filtering `_seq_no > seqNoFloor`, projecting only the version + // metadata columns, with pushdown enabled. + long streamPtr = executeInternalSearch( readerPtr, - substraitPlan, runtimePtr, - configSnapshot, + NativeBridge.INTERNAL_SEARCH_SEQ_NO_ABOVE, + seqNoFloor, "DataFusion range query failed" ); all.addAll(readAllRows(streamPtr)); @@ -202,41 +195,40 @@ private Map readSingleRow(long streamPtr) { } } - private long executeNativeQuery( - long readerPtr, - byte[] substraitPlan, - long runtimePtr, - WireConfigSnapshot configSnapshot, - String errorMessage - ) throws IOException { + /** + * Runs an engine-internal point lookup through {@link NativeBridge#executeQueryAsync} and + * returns the result stream pointer. No Substrait is generated: the native side builds the + * filter plan from {@code mode} + {@code bound} via the DataFrame API. {@code queryConfigPtr} + * is 0 — the internal-search path uses its own session config (pushdown on, single partition), + * so no {@link WireConfigSnapshot} is needed. + */ + private long executeInternalSearch(long readerPtr, long runtimePtr, long mode, long bound, String errorMessage) throws IOException { CompletableFuture future = new CompletableFuture<>(); - try (Arena arena = Arena.ofConfined()) { - MemorySegment configSegment = arena.allocate(WireConfigSnapshot.BYTE_SIZE); - configSnapshot.writeTo(configSegment); - NativeBridge.executeQueryAsync( - readerPtr, - GET_BY_ID_TABLE_ALIAS, - substraitPlan, - runtimePtr, - 0L, - configSegment.address(), - new ActionListener<>() { - @Override - public void onResponse(Long v) { - future.complete(v); - } + NativeBridge.executeQueryAsync( + readerPtr, + GET_BY_ID_TABLE_ALIAS, + EMPTY_PLAN, + runtimePtr, + 0L, + 0L, + mode, + bound, + new ActionListener<>() { + @Override + public void onResponse(Long v) { + future.complete(v); + } - @Override - public void onFailure(Exception e) { - future.completeExceptionally(e); - } + @Override + public void onFailure(Exception e) { + future.completeExceptionally(e); } - ); - try { - return future.join(); - } catch (Exception e) { - throw new IOException(errorMessage, e); } + ); + try { + return future.join(); + } catch (Exception e) { + throw new IOException(errorMessage, e); } } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java index 475f8d228191b..ac9e967022152 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java @@ -248,15 +248,17 @@ private static RuntimeException rethrowConverted(RuntimeException e) { EXECUTE_QUERY = linker.downcallHandle( lib.find("df_execute_query").orElseThrow(), FunctionDescriptor.of( - ValueLayout.JAVA_LONG, - ValueLayout.JAVA_LONG, - ValueLayout.ADDRESS, - ValueLayout.JAVA_LONG, - ValueLayout.ADDRESS, - ValueLayout.JAVA_LONG, - ValueLayout.JAVA_LONG, - ValueLayout.JAVA_LONG, - ValueLayout.JAVA_LONG + ValueLayout.JAVA_LONG, // returns stream_ptr + ValueLayout.JAVA_LONG, // shard_view_ptr + ValueLayout.ADDRESS, // table_name_ptr + ValueLayout.JAVA_LONG, // table_name_len + ValueLayout.ADDRESS, // plan_ptr + ValueLayout.JAVA_LONG, // plan_len + ValueLayout.JAVA_LONG, // runtime_ptr + ValueLayout.JAVA_LONG, // context_id + ValueLayout.JAVA_LONG, // query_config_ptr + ValueLayout.JAVA_LONG, // internal_search_mode + ValueLayout.JAVA_LONG // internal_search_bound ) ); @@ -901,6 +903,35 @@ public static void closeDatafusionReader(long ptr) { // ---- Query execution (confined Arena for tableName + plan bytes) ---- + /** {@code internal_search_mode}: normal query — decode {@code substraitPlan} as Substrait. */ + public static final long INTERNAL_SEARCH_OFF = 0L; + /** {@code internal_search_mode}: get-by-row-id — native plan filters {@code __row_id__ = bound}, {@code substraitPlan} ignored. */ + public static final long INTERNAL_SEARCH_BY_ROW_ID = 1L; + /** {@code internal_search_mode}: seq-no scan — native plan filters {@code _seq_no > bound}, {@code substraitPlan} ignored. */ + public static final long INTERNAL_SEARCH_SEQ_NO_ABOVE = 2L; + + public static void executeQueryAsync( + long readerPtr, + String tableName, + byte[] substraitPlan, + long runtimePtr, + long contextId, + long queryConfigPtr, + ActionListener listener + ) { + executeQueryAsync(readerPtr, tableName, substraitPlan, runtimePtr, contextId, queryConfigPtr, INTERNAL_SEARCH_OFF, 0L, listener); + } + + /** + * Executes a query and returns an opaque stream pointer via {@code listener}. + *

+ * When {@code internalSearchMode} is {@link #INTERNAL_SEARCH_OFF}, {@code substraitPlan} is + * decoded as a Substrait plan (normal search). When it is {@link #INTERNAL_SEARCH_BY_ROW_ID} + * or {@link #INTERNAL_SEARCH_SEQ_NO_ABOVE}, the native side ignores {@code substraitPlan} and + * builds a single pushed-down filter plan via the DataFusion DataFrame API, using + * {@code internalSearchBound} as the {@code __row_id__} value or the {@code _seq_no} floor. + * The returned stream is drained identically in all modes. + */ public static void executeQueryAsync( long readerPtr, String tableName, @@ -908,6 +939,8 @@ public static void executeQueryAsync( long runtimePtr, long contextId, long queryConfigPtr, + long internalSearchMode, + long internalSearchBound, ActionListener listener ) { try { @@ -928,7 +961,9 @@ public static void executeQueryAsync( (long) substraitPlan.length, runtimePtr, contextId, - queryConfigPtr + queryConfigPtr, + internalSearchMode, + internalSearchBound ); listener.onResponse(result); } catch (Throwable t) {