diff --git a/modules/store-subdirectory/src/main/java/org/opensearch/plugin/store/subdirectory/SubdirectoryAwareStore.java b/modules/store-subdirectory/src/main/java/org/opensearch/plugin/store/subdirectory/SubdirectoryAwareStore.java index 84fd9a5a97f34..46667f173b80e 100644 --- a/modules/store-subdirectory/src/main/java/org/opensearch/plugin/store/subdirectory/SubdirectoryAwareStore.java +++ b/modules/store-subdirectory/src/main/java/org/opensearch/plugin/store/subdirectory/SubdirectoryAwareStore.java @@ -16,10 +16,8 @@ import org.apache.lucene.index.SegmentInfos; import org.apache.lucene.store.Directory; import org.apache.lucene.store.FSDirectory; -import org.apache.lucene.store.FilterDirectory; import org.apache.lucene.store.IOContext; import org.apache.lucene.store.IndexInput; -import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.Version; import org.opensearch.common.lucene.Lucene; import org.opensearch.core.index.shard.ShardId; @@ -31,20 +29,12 @@ import org.opensearch.plugins.IndexStorePlugin; import java.io.IOException; -import java.nio.file.FileVisitResult; -import java.nio.file.Files; -import java.nio.file.NoSuchFileException; import java.nio.file.Path; -import java.nio.file.SimpleFileVisitor; -import java.nio.file.attribute.BasicFileAttributes; -import java.util.Arrays; -import java.util.Collection; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; import java.util.Map; import java.util.Set; -import java.util.stream.Collectors; /** * A store implementation that supports files organized in subdirectories. @@ -252,109 +242,17 @@ private void computeFileMetadata(String fileName, Map /** * A Lucene Directory implementation that handles files in subdirectories. - * - * This directory wrapper enables file operations across subdirectories within - * the shard data path. It resolves paths, creates necessary directory structures, - * and delegates actual file operations to appropriate filesystem locations. + * Extends the server's SubdirectoryAwareDirectory for backward compatibility. */ - public static class SubdirectoryAwareDirectory extends FilterDirectory { - private static final Set EXCLUDED_SUBDIRECTORIES = Set.of("index/", "translog/", "_state/"); - private final ShardPath shardPath; - + public static class SubdirectoryAwareDirectory extends org.opensearch.index.store.SubdirectoryAwareDirectory { /** - * Constructor for SubdirectoryAwareDirectory. + * Creates a new SubdirectoryAwareDirectory wrapping the given delegate. * - * @param delegate the delegate directory - * @param shardPath the shard path + * @param delegate the underlying Lucene directory + * @param shardPath the shard path for resolving subdirectories */ public SubdirectoryAwareDirectory(Directory delegate, ShardPath shardPath) { - super(delegate); - this.shardPath = shardPath; - } - - @Override - public IndexInput openInput(String name, IOContext context) throws IOException { - return super.openInput(parseFilePath(name), context); - } - - @Override - public IndexOutput createOutput(String name, IOContext context) throws IOException { - String targetFilePath = parseFilePath(name); - Path targetFile = Path.of(targetFilePath); - Files.createDirectories(targetFile.getParent()); - return super.createOutput(targetFilePath, context); - } - - @Override - public void deleteFile(String name) throws IOException { - super.deleteFile(parseFilePath(name)); - } - - @Override - public long fileLength(String name) throws IOException { - return super.fileLength(parseFilePath(name)); - } - - @Override - public void sync(Collection names) throws IOException { - super.sync(names.stream().map(this::parseFilePath).collect(Collectors.toList())); - } - - @Override - public void rename(String source, String dest) throws IOException { - super.rename(parseFilePath(source), parseFilePath(dest)); - } - - @Override - public String[] listAll() throws IOException { - // Get files from the delegate (regular index files) - String[] delegateFiles = super.listAll(); - - // Get subdirectory files by scanning all subdirectories - Set allFiles = new HashSet<>(Arrays.asList(delegateFiles)); - addSubdirectoryFiles(allFiles); - - return allFiles.stream().sorted().toArray(String[]::new); - } - - private void addSubdirectoryFiles(Set allFiles) throws IOException { - Path dataPath = shardPath.getDataPath(); - Files.walkFileTree(dataPath, new SimpleFileVisitor<>() { - @Override - public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) { - if (attrs.isRegularFile()) { - Path relativePath = dataPath.relativize(file); - // Only add files that are in subdirectories (have a parent directory) - if (relativePath.getParent() != null) { - String relativePathStr = relativePath.toString(); - // Exclude index dir (handled in super.listAll()), translog dir, and _state dir - if (EXCLUDED_SUBDIRECTORIES.stream().noneMatch(relativePathStr::startsWith)) { - allFiles.add(relativePathStr); - } - } - } - return FileVisitResult.CONTINUE; - } - - @Override - public FileVisitResult visitFileFailed(Path file, IOException e) throws IOException { - if (e instanceof NoSuchFileException) { - logger.debug("Skipping inaccessible file during size estimation: {}", file); - return FileVisitResult.CONTINUE; - } - throw e; - } - }); - } - - private String parseFilePath(String fileName) { - if (Path.of(fileName).getParent() != null) { - // File path (e.g., "subdirectory/segments_1" or "subdirectory/recovery.xxx.segments_1") - return shardPath.getDataPath().resolve(fileName).toString(); - } else { - // Simple filename (e.g., "segments_1") - resolve relative to the shard's index directory - return shardPath.resolveIndex().resolve(fileName).toString(); - } + super(delegate, shardPath); } } } diff --git a/sandbox/libs/dataformat-native/rust/Cargo.toml b/sandbox/libs/dataformat-native/rust/Cargo.toml index b34cd03bd2d40..c69ed2fa6c9b5 100644 --- a/sandbox/libs/dataformat-native/rust/Cargo.toml +++ b/sandbox/libs/dataformat-native/rust/Cargo.toml @@ -48,6 +48,7 @@ url = "2.0" tempfile = "3.0" chrono = "0.4" once_cell = "1.21.3" +crc32fast = "1.4" parking_lot = "0.12.5" lazy_static = "1.4.0" criterion = { version = "0.5", features = ["async_tokio"] } diff --git a/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneReaderManagerTests.java b/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneReaderManagerTests.java index 591ee1952aab1..7d2eec3cfae3a 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneReaderManagerTests.java +++ b/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneReaderManagerTests.java @@ -134,12 +134,27 @@ public CatalogSnapshot cloneNoAcquire() { } @Override - public void setUserData(Map userData) {} + public void setUserData(Map userData, boolean commitData) {} @Override public CatalogSnapshot clone() { return this; } + + @Override + public int getFormatVersionForFile(String file) { + return 0; + } + + @Override + public byte[] serialize() throws IOException { + return new byte[0]; + } + + @Override + public Collection getFiles(boolean includeSegmentsFile) { + return List.of(); + } }; } diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/DefaultPlanExecutorTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/DefaultPlanExecutorTests.java index d235f6fe4a3e4..6f6eb7d113ca5 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/DefaultPlanExecutorTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/DefaultPlanExecutorTests.java @@ -292,13 +292,28 @@ public String serializeToString() { } @Override - public void setUserData(Map userData) {} + public void setUserData(Map userData, boolean commitData) {} @Override public MockCatalogSnapshot clone() { return new MockCatalogSnapshot(generation, segments, format); } + @Override + public int getFormatVersionForFile(String file) { + return 0; + } + + @Override + public byte[] serialize() throws IOException { + return new byte[0]; + } + + @Override + public Collection getFiles(boolean includeSegmentsFile) { + return List.of(); + } + @Override protected void closeInternal() {} } diff --git a/sandbox/plugins/composite-engine/README.md b/sandbox/plugins/composite-engine/README.md index 82a8fdfb44010..8a4bd2aa5057e 100644 --- a/sandbox/plugins/composite-engine/README.md +++ b/sandbox/plugins/composite-engine/README.md @@ -23,7 +23,7 @@ Format plugins (e.g., Parquet) extend this plugin by declaring `extendedPlugins ## Key classes -- **`CompositeEnginePlugin`** — The `ExtensiblePlugin` entry point. Discovers format plugins, validates settings, and creates the composite engine. +- **`CompositeDataFormatPlugin`** — The `ExtensiblePlugin` entry point. Discovers format plugins, validates settings, and creates the composite engine. - **`CompositeIndexingExecutionEngine`** — Orchestrates indexing across primary and secondary format engines. - **`CompositeDataFormat`** — A `DataFormat` that wraps multiple per-format instances. - **`CompositeDocumentInput`** — Routes field additions to the appropriate per-format `DocumentInput` based on field type capabilities. diff --git a/sandbox/plugins/composite-engine/build.gradle b/sandbox/plugins/composite-engine/build.gradle index 1547c0c33f445..92c4ed490d19d 100644 --- a/sandbox/plugins/composite-engine/build.gradle +++ b/sandbox/plugins/composite-engine/build.gradle @@ -8,7 +8,7 @@ opensearchplugin { description = 'Composite indexing engine plugin that orchestrates multi-format indexing across multiple data format engines.' - classname = 'org.opensearch.composite.CompositeEnginePlugin' + classname = 'org.opensearch.composite.CompositeDataFormatPlugin' } dependencies { diff --git a/sandbox/plugins/composite-engine/src/main/java/org/opensearch/composite/CompositeEnginePlugin.java b/sandbox/plugins/composite-engine/src/main/java/org/opensearch/composite/CompositeDataFormatPlugin.java similarity index 76% rename from sandbox/plugins/composite-engine/src/main/java/org/opensearch/composite/CompositeEnginePlugin.java rename to sandbox/plugins/composite-engine/src/main/java/org/opensearch/composite/CompositeDataFormatPlugin.java index 3b015a75f4629..612711b85117f 100644 --- a/sandbox/plugins/composite-engine/src/main/java/org/opensearch/composite/CompositeEnginePlugin.java +++ b/sandbox/plugins/composite-engine/src/main/java/org/opensearch/composite/CompositeDataFormatPlugin.java @@ -12,10 +12,14 @@ import org.apache.logging.log4j.Logger; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; +import org.opensearch.index.IndexSettings; import org.opensearch.index.engine.dataformat.DataFormat; +import org.opensearch.index.engine.dataformat.DataFormatDescriptor; import org.opensearch.index.engine.dataformat.DataFormatPlugin; import org.opensearch.index.engine.dataformat.IndexingEngineConfig; import org.opensearch.index.engine.dataformat.IndexingExecutionEngine; +import org.opensearch.index.store.FormatChecksumStrategy; import org.opensearch.plugins.ExtensiblePlugin; import org.opensearch.plugins.Plugin; @@ -44,9 +48,9 @@ * @opensearch.experimental */ @ExperimentalApi -public class CompositeEnginePlugin extends Plugin implements ExtensiblePlugin, DataFormatPlugin { +public class CompositeDataFormatPlugin extends Plugin implements ExtensiblePlugin, DataFormatPlugin { - private static final Logger logger = LogManager.getLogger(CompositeEnginePlugin.class); + private static final Logger logger = LogManager.getLogger(CompositeDataFormatPlugin.class); /** * Index setting that designates the primary data format for an index. @@ -78,9 +82,10 @@ public class CompositeEnginePlugin extends Plugin implements ExtensiblePlugin, D * {@link DataFormat#priority()} is retained. */ private volatile Map dataFormatPlugins = Map.of(); + private volatile Map lastDescriptors = Map.of(); /** Creates a new composite engine plugin. */ - public CompositeEnginePlugin() {} + public CompositeDataFormatPlugin() {} @Override public void loadExtensions(ExtensionLoader loader) { @@ -135,16 +140,42 @@ public DataFormat getDataFormat() { } @Override - public IndexingExecutionEngine indexingEngine(IndexingEngineConfig settings) { + public IndexingExecutionEngine indexingEngine(IndexingEngineConfig settings, FormatChecksumStrategy checksumStrategy) { + Map strategies = new HashMap<>(); + for (Map.Entry entry : lastDescriptors.entrySet()) { + strategies.put(entry.getKey(), entry.getValue().getChecksumStrategy()); + } return new CompositeIndexingExecutionEngine( dataFormatPlugins, settings.indexSettings(), settings.mapperService(), settings.shardPath(), - settings.committer() + settings.committer(), + strategies ); } + @Override + public Map getFormatDescriptors(IndexSettings indexSettings) { + Settings settings = indexSettings.getSettings(); + String primaryFormatName = PRIMARY_DATA_FORMAT.get(settings); + List secondaryFormatNames = SECONDARY_DATA_FORMATS.get(settings); + + Map descriptors = new HashMap<>(); + DataFormatPlugin primaryPlugin = dataFormatPlugins.get(primaryFormatName); + if (primaryPlugin != null) { + descriptors.putAll(primaryPlugin.getFormatDescriptors(indexSettings)); + } + for (String secondaryName : secondaryFormatNames) { + DataFormatPlugin secondaryPlugin = dataFormatPlugins.get(secondaryName); + if (secondaryPlugin != null) { + descriptors.putAll(secondaryPlugin.getFormatDescriptors(indexSettings)); + } + } + lastDescriptors = Map.copyOf(descriptors); + return lastDescriptors; + } + /** * Returns the discovered data format plugins keyed by format name. * diff --git a/sandbox/plugins/composite-engine/src/main/java/org/opensearch/composite/CompositeIndexingExecutionEngine.java b/sandbox/plugins/composite-engine/src/main/java/org/opensearch/composite/CompositeIndexingExecutionEngine.java index 79c7f416c0674..c82dd374452a7 100644 --- a/sandbox/plugins/composite-engine/src/main/java/org/opensearch/composite/CompositeIndexingExecutionEngine.java +++ b/sandbox/plugins/composite-engine/src/main/java/org/opensearch/composite/CompositeIndexingExecutionEngine.java @@ -33,6 +33,7 @@ import org.opensearch.index.engine.exec.coord.CatalogSnapshotManager; import org.opensearch.index.mapper.MapperService; import org.opensearch.index.shard.ShardPath; +import org.opensearch.index.store.FormatChecksumStrategy; import java.io.IOException; import java.util.ArrayList; @@ -85,11 +86,13 @@ public class CompositeIndexingExecutionEngine implements IndexingExecutionEngine * The writer pool is created internally and initialized with a writer supplier * that creates {@link CompositeWriter} instances bound to this engine. * - * @param dataFormatPlugins the discovered data format plugins keyed by format name - * @param indexSettings the index settings containing composite configuration - * @param mapperService the mapper service for field mapping resolution - * @param shardPath the shard path for file storage - * @param committer the committer for durable catalog snapshot persistence during flush + * @param dataFormatPlugins the discovered data format plugins keyed by format name + * @param indexSettings the index settings containing composite configuration + * @param mapperService the mapper service for field mapping resolution + * @param shardPath the shard path for file storage + * @param committer the committer for durable catalog snapshot persistence during flush + * @param checksumStrategies per-format checksum strategies from the directory, keyed by format name. + * May be null or empty if the directory is not yet available. * @throws IllegalArgumentException if any configured format is not registered * @throws IllegalStateException if committer is null */ @@ -98,7 +101,8 @@ public CompositeIndexingExecutionEngine( IndexSettings indexSettings, MapperService mapperService, ShardPath shardPath, - Committer committer + Committer committer, + Map checksumStrategies ) { Objects.requireNonNull(dataFormatPlugins, "dataFormatPlugins must not be null"); Objects.requireNonNull(indexSettings, "indexSettings must not be null"); @@ -108,22 +112,23 @@ public CompositeIndexingExecutionEngine( Settings settings = indexSettings.getSettings(); - String primaryFormatName = CompositeEnginePlugin.PRIMARY_DATA_FORMAT.get(settings); - List secondaryFormatNames = CompositeEnginePlugin.SECONDARY_DATA_FORMATS.get(settings); + String primaryFormatName = CompositeDataFormatPlugin.PRIMARY_DATA_FORMAT.get(settings); + List secondaryFormatNames = CompositeDataFormatPlugin.SECONDARY_DATA_FORMATS.get(settings); validateFormatsRegistered(dataFormatPlugins, primaryFormatName, secondaryFormatNames); + Map strategies = checksumStrategies != null ? checksumStrategies : Map.of(); IndexingEngineConfig engineSettings = new IndexingEngineConfig(committer, mapperService, shardPath, indexSettings, null); List allFormats = new ArrayList<>(); DataFormatPlugin primaryPlugin = dataFormatPlugins.get(primaryFormatName); - this.primaryEngine = primaryPlugin.indexingEngine(engineSettings); + this.primaryEngine = primaryPlugin.indexingEngine(engineSettings, strategies.get(primaryFormatName)); allFormats.add(primaryPlugin.getDataFormat()); List> secondaries = new ArrayList<>(); for (String secondaryName : secondaryFormatNames) { DataFormatPlugin secondaryPlugin = dataFormatPlugins.get(secondaryName); - secondaries.add(secondaryPlugin.indexingEngine(engineSettings)); + secondaries.add(secondaryPlugin.indexingEngine(engineSettings, strategies.get(secondaryName))); allFormats.add(secondaryPlugin.getDataFormat()); } this.secondaryEngines = Set.copyOf(secondaries); diff --git a/sandbox/plugins/composite-engine/src/test/java/org/opensearch/composite/CompositeDataFormatPluginTests.java b/sandbox/plugins/composite-engine/src/test/java/org/opensearch/composite/CompositeDataFormatPluginTests.java new file mode 100644 index 0000000000000..f9572d40a9a60 --- /dev/null +++ b/sandbox/plugins/composite-engine/src/test/java/org/opensearch/composite/CompositeDataFormatPluginTests.java @@ -0,0 +1,239 @@ +/* + * 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.composite; + +import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; +import org.opensearch.index.IndexSettings; +import org.opensearch.index.engine.dataformat.DataFormat; +import org.opensearch.index.engine.dataformat.DataFormatPlugin; +import org.opensearch.plugins.ExtensiblePlugin; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Collections; +import java.util.List; +import java.util.Map; + +/** + * Tests for {@link CompositeDataFormatPlugin}. + */ +public class CompositeDataFormatPluginTests extends OpenSearchTestCase { + + public void testGetSettingsReturnsBothSettings() { + CompositeDataFormatPlugin plugin = new CompositeDataFormatPlugin(); + List> settings = plugin.getSettings(); + assertEquals(2, settings.size()); + assertTrue(settings.contains(CompositeDataFormatPlugin.PRIMARY_DATA_FORMAT)); + assertTrue(settings.contains(CompositeDataFormatPlugin.SECONDARY_DATA_FORMATS)); + } + + public void testPrimaryDataFormatDefaultsToLucene() { + Settings settings = Settings.builder().build(); + assertEquals("lucene", CompositeDataFormatPlugin.PRIMARY_DATA_FORMAT.get(settings)); + } + + public void testSecondaryDataFormatsDefaultsToEmpty() { + Settings settings = Settings.builder().build(); + assertTrue(CompositeDataFormatPlugin.SECONDARY_DATA_FORMATS.get(settings).isEmpty()); + } + + public void testGetDataFormatReturnsNull() { + CompositeDataFormatPlugin plugin = new CompositeDataFormatPlugin(); + assertNull(plugin.getDataFormat()); + } + + public void testLoadExtensionsRegistersPlugins() { + CompositeDataFormatPlugin plugin = new CompositeDataFormatPlugin(); + DataFormatPlugin lucenePlugin = CompositeTestHelper.stubPlugin("lucene", 1); + DataFormatPlugin parquetPlugin = CompositeTestHelper.stubPlugin("parquet", 2); + + plugin.loadExtensions(new ExtensiblePlugin.ExtensionLoader() { + @Override + @SuppressWarnings("unchecked") + public List loadExtensions(Class extensionPointType) { + if (extensionPointType == DataFormatPlugin.class) { + return (List) List.of(lucenePlugin, parquetPlugin); + } + return Collections.emptyList(); + } + }); + + Map plugins = plugin.getDataFormatPlugins(); + assertEquals(2, plugins.size()); + assertTrue(plugins.containsKey("lucene")); + assertTrue(plugins.containsKey("parquet")); + } + + public void testLoadExtensionsHigherPriorityWins() { + CompositeDataFormatPlugin plugin = new CompositeDataFormatPlugin(); + DataFormatPlugin lowPriority = CompositeTestHelper.stubPlugin("lucene", 1); + DataFormatPlugin highPriority = CompositeTestHelper.stubPlugin("lucene", 100); + + plugin.loadExtensions(new ExtensiblePlugin.ExtensionLoader() { + @Override + @SuppressWarnings("unchecked") + public List loadExtensions(Class extensionPointType) { + if (extensionPointType == DataFormatPlugin.class) { + return (List) List.of(lowPriority, highPriority); + } + return Collections.emptyList(); + } + }); + + Map plugins = plugin.getDataFormatPlugins(); + assertEquals(1, plugins.size()); + // The high priority one should win + assertEquals(100, plugins.get("lucene").getDataFormat().priority()); + } + + public void testLoadExtensionsSkipsNullDataFormat() { + CompositeDataFormatPlugin plugin = new CompositeDataFormatPlugin(); + DataFormatPlugin nullPlugin = new DataFormatPlugin() { + @Override + public DataFormat getDataFormat() { + return null; + } + + @Override + public org.opensearch.index.engine.dataformat.IndexingExecutionEngine indexingEngine( + org.opensearch.index.engine.dataformat.IndexingEngineConfig settings, + org.opensearch.index.store.FormatChecksumStrategy checksumStrategy + ) { + return null; + } + }; + + plugin.loadExtensions(new ExtensiblePlugin.ExtensionLoader() { + @Override + @SuppressWarnings("unchecked") + public List loadExtensions(Class extensionPointType) { + if (extensionPointType == DataFormatPlugin.class) { + return (List) List.of(nullPlugin); + } + return Collections.emptyList(); + } + }); + + assertTrue(plugin.getDataFormatPlugins().isEmpty()); + } + + public void testLoadExtensionsWithEmptyList() { + CompositeDataFormatPlugin plugin = new CompositeDataFormatPlugin(); + plugin.loadExtensions(new ExtensiblePlugin.ExtensionLoader() { + @Override + public List loadExtensions(Class extensionPointType) { + return Collections.emptyList(); + } + }); + + assertTrue(plugin.getDataFormatPlugins().isEmpty()); + } + + public void testGetFormatDescriptorsDelegatestoPlugins() { + CompositeDataFormatPlugin plugin = new CompositeDataFormatPlugin(); + + // Create a plugin that returns a descriptor + DataFormatPlugin parquetPlugin = new DataFormatPlugin() { + @Override + public DataFormat getDataFormat() { + return CompositeTestHelper.stubFormat("parquet", 2, java.util.Set.of()); + } + + @Override + public org.opensearch.index.engine.dataformat.IndexingExecutionEngine indexingEngine( + org.opensearch.index.engine.dataformat.IndexingEngineConfig settings, + org.opensearch.index.store.FormatChecksumStrategy checksumStrategy + ) { + return null; + } + + @Override + public Map getFormatDescriptors( + IndexSettings indexSettings + ) { + return Map.of( + "parquet", + new org.opensearch.index.engine.dataformat.DataFormatDescriptor( + "parquet", + new org.opensearch.index.store.checksum.GenericCRC32ChecksumHandler() + ) + ); + } + }; + + plugin.loadExtensions(new ExtensiblePlugin.ExtensionLoader() { + @Override + @SuppressWarnings("unchecked") + public List loadExtensions(Class extensionPointType) { + if (extensionPointType == DataFormatPlugin.class) { + return (List) List.of(parquetPlugin); + } + return Collections.emptyList(); + } + }); + + // Build index settings with parquet as secondary + Settings settings = Settings.builder() + .put("index.composite.primary_data_format", "lucene") + .putList("index.composite.secondary_data_formats", "parquet") + .put(org.opensearch.cluster.metadata.IndexMetadata.SETTING_VERSION_CREATED, org.opensearch.Version.CURRENT) + .put(org.opensearch.cluster.metadata.IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0) + .put(org.opensearch.cluster.metadata.IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1) + .build(); + org.opensearch.cluster.metadata.IndexMetadata indexMetadata = org.opensearch.cluster.metadata.IndexMetadata.builder("test-index") + .settings(settings) + .build(); + IndexSettings indexSettings = new IndexSettings(indexMetadata, Settings.EMPTY); + + Map descriptors = plugin.getFormatDescriptors(indexSettings); + assertEquals(1, descriptors.size()); + assertTrue(descriptors.containsKey("parquet")); + assertEquals("parquet", descriptors.get("parquet").getFormatName()); + } + + public void testGetFormatDescriptorsEmptyWhenNoPluginsMatch() { + CompositeDataFormatPlugin plugin = new CompositeDataFormatPlugin(); + plugin.loadExtensions(new ExtensiblePlugin.ExtensionLoader() { + @Override + public List loadExtensions(Class extensionPointType) { + return Collections.emptyList(); + } + }); + + Settings settings = Settings.builder() + .put(org.opensearch.cluster.metadata.IndexMetadata.SETTING_VERSION_CREATED, org.opensearch.Version.CURRENT) + .put(org.opensearch.cluster.metadata.IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0) + .put(org.opensearch.cluster.metadata.IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1) + .build(); + org.opensearch.cluster.metadata.IndexMetadata indexMetadata = org.opensearch.cluster.metadata.IndexMetadata.builder("test-index") + .settings(settings) + .build(); + IndexSettings indexSettings = new IndexSettings(indexMetadata, Settings.EMPTY); + + Map descriptors = plugin.getFormatDescriptors(indexSettings); + assertTrue(descriptors.isEmpty()); + } + + public void testGetDataFormatPluginsReturnsUnmodifiableMap() { + CompositeDataFormatPlugin plugin = new CompositeDataFormatPlugin(); + plugin.loadExtensions(new ExtensiblePlugin.ExtensionLoader() { + @Override + @SuppressWarnings("unchecked") + public List loadExtensions(Class extensionPointType) { + if (extensionPointType == DataFormatPlugin.class) { + return (List) List.of(CompositeTestHelper.stubPlugin("lucene", 1)); + } + return Collections.emptyList(); + } + }); + + Map plugins = plugin.getDataFormatPlugins(); + expectThrows(UnsupportedOperationException.class, () -> plugins.put("new", CompositeTestHelper.stubPlugin("new", 1))); + } +} diff --git a/sandbox/plugins/composite-engine/src/test/java/org/opensearch/composite/CompositeEnginePluginTests.java b/sandbox/plugins/composite-engine/src/test/java/org/opensearch/composite/CompositeEnginePluginTests.java deleted file mode 100644 index ab255d2737faa..0000000000000 --- a/sandbox/plugins/composite-engine/src/test/java/org/opensearch/composite/CompositeEnginePluginTests.java +++ /dev/null @@ -1,151 +0,0 @@ -/* - * 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.composite; - -import org.opensearch.common.settings.Setting; -import org.opensearch.common.settings.Settings; -import org.opensearch.index.engine.dataformat.DataFormat; -import org.opensearch.index.engine.dataformat.DataFormatPlugin; -import org.opensearch.index.engine.dataformat.IndexingEngineConfig; -import org.opensearch.plugins.ExtensiblePlugin; -import org.opensearch.test.OpenSearchTestCase; - -import java.util.Collections; -import java.util.List; -import java.util.Map; - -/** - * Tests for {@link CompositeEnginePlugin}. - */ -public class CompositeEnginePluginTests extends OpenSearchTestCase { - - public void testGetSettingsReturnsBothSettings() { - CompositeEnginePlugin plugin = new CompositeEnginePlugin(); - List> settings = plugin.getSettings(); - assertEquals(2, settings.size()); - assertTrue(settings.contains(CompositeEnginePlugin.PRIMARY_DATA_FORMAT)); - assertTrue(settings.contains(CompositeEnginePlugin.SECONDARY_DATA_FORMATS)); - } - - public void testPrimaryDataFormatDefaultsToLucene() { - Settings settings = Settings.builder().build(); - assertEquals("lucene", CompositeEnginePlugin.PRIMARY_DATA_FORMAT.get(settings)); - } - - public void testSecondaryDataFormatsDefaultsToEmpty() { - Settings settings = Settings.builder().build(); - assertTrue(CompositeEnginePlugin.SECONDARY_DATA_FORMATS.get(settings).isEmpty()); - } - - public void testGetDataFormatReturnsNull() { - CompositeEnginePlugin plugin = new CompositeEnginePlugin(); - assertNull(plugin.getDataFormat()); - } - - public void testLoadExtensionsRegistersPlugins() { - CompositeEnginePlugin plugin = new CompositeEnginePlugin(); - DataFormatPlugin lucenePlugin = CompositeTestHelper.stubPlugin("lucene", 1); - DataFormatPlugin parquetPlugin = CompositeTestHelper.stubPlugin("parquet", 2); - - plugin.loadExtensions(new ExtensiblePlugin.ExtensionLoader() { - @Override - @SuppressWarnings("unchecked") - public List loadExtensions(Class extensionPointType) { - if (extensionPointType == DataFormatPlugin.class) { - return (List) List.of(lucenePlugin, parquetPlugin); - } - return Collections.emptyList(); - } - }); - - Map plugins = plugin.getDataFormatPlugins(); - assertEquals(2, plugins.size()); - assertTrue(plugins.containsKey("lucene")); - assertTrue(plugins.containsKey("parquet")); - } - - public void testLoadExtensionsHigherPriorityWins() { - CompositeEnginePlugin plugin = new CompositeEnginePlugin(); - DataFormatPlugin lowPriority = CompositeTestHelper.stubPlugin("lucene", 1); - DataFormatPlugin highPriority = CompositeTestHelper.stubPlugin("lucene", 100); - - plugin.loadExtensions(new ExtensiblePlugin.ExtensionLoader() { - @Override - @SuppressWarnings("unchecked") - public List loadExtensions(Class extensionPointType) { - if (extensionPointType == DataFormatPlugin.class) { - return (List) List.of(lowPriority, highPriority); - } - return Collections.emptyList(); - } - }); - - Map plugins = plugin.getDataFormatPlugins(); - assertEquals(1, plugins.size()); - // The high priority one should win - assertEquals(100, plugins.get("lucene").getDataFormat().priority()); - } - - public void testLoadExtensionsSkipsNullDataFormat() { - CompositeEnginePlugin plugin = new CompositeEnginePlugin(); - DataFormatPlugin nullPlugin = new DataFormatPlugin() { - @Override - public DataFormat getDataFormat() { - return null; - } - - @Override - public org.opensearch.index.engine.dataformat.IndexingExecutionEngine indexingEngine(IndexingEngineConfig settings) { - return null; - } - }; - - plugin.loadExtensions(new ExtensiblePlugin.ExtensionLoader() { - @Override - @SuppressWarnings("unchecked") - public List loadExtensions(Class extensionPointType) { - if (extensionPointType == DataFormatPlugin.class) { - return (List) List.of(nullPlugin); - } - return Collections.emptyList(); - } - }); - - assertTrue(plugin.getDataFormatPlugins().isEmpty()); - } - - public void testLoadExtensionsWithEmptyList() { - CompositeEnginePlugin plugin = new CompositeEnginePlugin(); - plugin.loadExtensions(new ExtensiblePlugin.ExtensionLoader() { - @Override - public List loadExtensions(Class extensionPointType) { - return Collections.emptyList(); - } - }); - - assertTrue(plugin.getDataFormatPlugins().isEmpty()); - } - - public void testGetDataFormatPluginsReturnsUnmodifiableMap() { - CompositeEnginePlugin plugin = new CompositeEnginePlugin(); - plugin.loadExtensions(new ExtensiblePlugin.ExtensionLoader() { - @Override - @SuppressWarnings("unchecked") - public List loadExtensions(Class extensionPointType) { - if (extensionPointType == DataFormatPlugin.class) { - return (List) List.of(CompositeTestHelper.stubPlugin("lucene", 1)); - } - return Collections.emptyList(); - } - }); - - Map plugins = plugin.getDataFormatPlugins(); - expectThrows(UnsupportedOperationException.class, () -> plugins.put("new", CompositeTestHelper.stubPlugin("new", 1))); - } -} diff --git a/sandbox/plugins/composite-engine/src/test/java/org/opensearch/composite/CompositeIndexingExecutionEngineTests.java b/sandbox/plugins/composite-engine/src/test/java/org/opensearch/composite/CompositeIndexingExecutionEngineTests.java index 5618be857f31a..7a59c7fd58880 100644 --- a/sandbox/plugins/composite-engine/src/test/java/org/opensearch/composite/CompositeIndexingExecutionEngineTests.java +++ b/sandbox/plugins/composite-engine/src/test/java/org/opensearch/composite/CompositeIndexingExecutionEngineTests.java @@ -56,7 +56,7 @@ public void testConstructorThrowsWhenPrimaryFormatNotRegistered() { IndexSettings indexSettings = createIndexSettings("parquet"); IllegalArgumentException ex = expectThrows( IllegalArgumentException.class, - () -> new CompositeIndexingExecutionEngine(plugins, indexSettings, null, null, new CompositeTestHelper.StubCommitter()) + () -> new CompositeIndexingExecutionEngine(plugins, indexSettings, null, null, new CompositeTestHelper.StubCommitter(), null) ); assertTrue(ex.getMessage().contains("parquet")); } @@ -77,7 +77,7 @@ public void testConstructorThrowsWhenSecondaryFormatNotRegistered() { IllegalArgumentException ex = expectThrows( IllegalArgumentException.class, - () -> new CompositeIndexingExecutionEngine(plugins, indexSettings, null, null, new CompositeTestHelper.StubCommitter()) + () -> new CompositeIndexingExecutionEngine(plugins, indexSettings, null, null, new CompositeTestHelper.StubCommitter(), null) ); assertTrue(ex.getMessage().contains("parquet")); } @@ -86,7 +86,7 @@ public void testConstructorRejectsNullDataFormatPlugins() { IndexSettings indexSettings = createIndexSettings("lucene"); expectThrows( NullPointerException.class, - () -> new CompositeIndexingExecutionEngine(null, indexSettings, null, null, new CompositeTestHelper.StubCommitter()) + () -> new CompositeIndexingExecutionEngine(null, indexSettings, null, null, new CompositeTestHelper.StubCommitter(), null) ); } @@ -94,7 +94,7 @@ public void testConstructorRejectsNullIndexSettings() { Map plugins = Map.of("lucene", CompositeTestHelper.stubPlugin("lucene", 1)); expectThrows( NullPointerException.class, - () -> new CompositeIndexingExecutionEngine(plugins, null, null, null, new CompositeTestHelper.StubCommitter()) + () -> new CompositeIndexingExecutionEngine(plugins, null, null, null, new CompositeTestHelper.StubCommitter(), null) ); } @@ -195,7 +195,7 @@ public void testConstructorThrowsWhenCommitterNull() { IllegalStateException ex = expectThrows( IllegalStateException.class, - () -> new CompositeIndexingExecutionEngine(plugins, indexSettings, null, null, null) + () -> new CompositeIndexingExecutionEngine(plugins, indexSettings, null, null, null, null) ); assertTrue(ex.getMessage().contains("Committer must not be null")); } @@ -214,7 +214,7 @@ public void testRefreshNeverCallsCommitterMethods() throws IOException { plugins.put("lucene", CompositeTestHelper.stubPlugin("lucene", 1)); IndexSettings indexSettings = createIndexSettings("lucene"); - CompositeIndexingExecutionEngine engine = new CompositeIndexingExecutionEngine(plugins, indexSettings, null, null, tracking); + CompositeIndexingExecutionEngine engine = new CompositeIndexingExecutionEngine(plugins, indexSettings, null, null, tracking, null); // Reset tracking after construction (init is called during construction) tracking.commitCalled = false; @@ -235,7 +235,7 @@ public void testInitCalledDuringConstruction() { plugins.put("lucene", CompositeTestHelper.stubPlugin("lucene", 1)); IndexSettings indexSettings = createIndexSettings("lucene"); - CompositeIndexingExecutionEngine engine = new CompositeIndexingExecutionEngine(plugins, indexSettings, null, null, stub); + CompositeIndexingExecutionEngine engine = new CompositeIndexingExecutionEngine(plugins, indexSettings, null, null, stub, null); assertNotNull(engine); } @@ -245,7 +245,7 @@ public void testCloseCalledDuringShutdown() { plugins.put("lucene", CompositeTestHelper.stubPlugin("lucene", 1)); IndexSettings indexSettings = createIndexSettings("lucene"); - CompositeIndexingExecutionEngine engine = new CompositeIndexingExecutionEngine(plugins, indexSettings, null, null, stub); + CompositeIndexingExecutionEngine engine = new CompositeIndexingExecutionEngine(plugins, indexSettings, null, null, stub, null); engine.close(); assertTrue("close() must be called during shutdown", stub.closeCalled); } @@ -285,7 +285,14 @@ public SafeCommitInfo getSafeCommitInfo() { plugins.put("lucene", CompositeTestHelper.stubPlugin("lucene", 1)); IndexSettings indexSettings = createIndexSettings("lucene"); - CompositeIndexingExecutionEngine engine = new CompositeIndexingExecutionEngine(plugins, indexSettings, null, null, failingClose); + CompositeIndexingExecutionEngine engine = new CompositeIndexingExecutionEngine( + plugins, + indexSettings, + null, + null, + failingClose, + null + ); // close() should not throw — it logs the error and continues engine.close(); @@ -297,7 +304,7 @@ public void testFlushCallsCommitterCommit() throws IOException { plugins.put("lucene", CompositeTestHelper.stubPlugin("lucene", 1)); IndexSettings indexSettings = createIndexSettings("lucene"); - CompositeIndexingExecutionEngine engine = new CompositeIndexingExecutionEngine(plugins, indexSettings, null, null, tracking); + CompositeIndexingExecutionEngine engine = new CompositeIndexingExecutionEngine(plugins, indexSettings, null, null, tracking, null); CatalogSnapshotManager csm = new CatalogSnapshotManager(0, 0, 0, List.of(), 0, Map.of()); engine.setCatalogSnapshotManager(csm); @@ -337,7 +344,14 @@ public SafeCommitInfo getSafeCommitInfo() { plugins.put("lucene", CompositeTestHelper.stubPlugin("lucene", 1)); IndexSettings indexSettings = createIndexSettings("lucene"); - CompositeIndexingExecutionEngine engine = new CompositeIndexingExecutionEngine(plugins, indexSettings, null, null, failingCommit); + CompositeIndexingExecutionEngine engine = new CompositeIndexingExecutionEngine( + plugins, + indexSettings, + null, + null, + failingCommit, + null + ); CatalogSnapshotManager csm = new CatalogSnapshotManager(0, 0, 0, List.of(), 0, Map.of()); engine.setCatalogSnapshotManager(csm); diff --git a/sandbox/plugins/composite-engine/src/test/java/org/opensearch/composite/CompositeTestHelper.java b/sandbox/plugins/composite-engine/src/test/java/org/opensearch/composite/CompositeTestHelper.java index e07c125a89f94..a2048d4293174 100644 --- a/sandbox/plugins/composite-engine/src/test/java/org/opensearch/composite/CompositeTestHelper.java +++ b/sandbox/plugins/composite-engine/src/test/java/org/opensearch/composite/CompositeTestHelper.java @@ -28,6 +28,7 @@ import org.opensearch.index.engine.dataformat.Writer; import org.opensearch.index.engine.exec.commit.Committer; import org.opensearch.index.engine.exec.commit.IndexStoreProvider; +import org.opensearch.index.store.FormatChecksumStrategy; import java.util.Collection; import java.util.Collections; @@ -67,7 +68,7 @@ static CompositeIndexingExecutionEngine createStubEngine(String primaryName, Str IndexMetadata indexMetadata = IndexMetadata.builder("test-index").settings(settings).build(); IndexSettings indexSettings = new IndexSettings(indexMetadata, Settings.EMPTY); - return new CompositeIndexingExecutionEngine(plugins, indexSettings, null, null, new StubCommitter()); + return new CompositeIndexingExecutionEngine(plugins, indexSettings, null, null, new StubCommitter(), null); } static DataFormatPlugin stubPlugin(String formatName, long priority) { @@ -79,7 +80,7 @@ public DataFormat getDataFormat() { } @Override - public IndexingExecutionEngine indexingEngine(IndexingEngineConfig settings) { + public IndexingExecutionEngine indexingEngine(IndexingEngineConfig settings, FormatChecksumStrategy checksumStrategy) { return new StubIndexingExecutionEngine(format); } }; @@ -94,7 +95,7 @@ public DataFormat getDataFormat() { } @Override - public IndexingExecutionEngine indexingEngine(IndexingEngineConfig settings) { + public IndexingExecutionEngine indexingEngine(IndexingEngineConfig settings, FormatChecksumStrategy checksumStrategy) { return new StubIndexingExecutionEngine(format); } }; diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetDataFormatPlugin.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetDataFormatPlugin.java index a89db49cd94df..271b412d3822b 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetDataFormatPlugin.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetDataFormatPlugin.java @@ -17,10 +17,14 @@ 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.engine.dataformat.DataFormat; +import org.opensearch.index.engine.dataformat.DataFormatDescriptor; import org.opensearch.index.engine.dataformat.DataFormatPlugin; import org.opensearch.index.engine.dataformat.IndexingEngineConfig; import org.opensearch.index.engine.dataformat.IndexingExecutionEngine; +import org.opensearch.index.store.FormatChecksumStrategy; +import org.opensearch.index.store.PrecomputedChecksumStrategy; import org.opensearch.parquet.engine.ParquetDataFormat; import org.opensearch.parquet.engine.ParquetIndexingEngine; import org.opensearch.parquet.fields.ArrowSchemaBuilder; @@ -36,6 +40,7 @@ import java.util.Collection; import java.util.Collections; import java.util.List; +import java.util.Map; import java.util.function.Supplier; /** @@ -46,6 +51,13 @@ * {@link #createComponents} and passes them to the per-shard * {@link ParquetIndexingEngine} instances created in {@link #indexingEngine}. * + *

The descriptor provides a {@link PrecomputedChecksumStrategy} that the directory + * holds at construction time. The {@link ParquetIndexingEngine} receives the same + * strategy instance from the directory via + * {@link org.opensearch.index.store.DataFormatAwareStoreDirectory#getChecksumStrategy}, + * so pre-computed CRC32 values registered during write are directly visible to the + * upload path — no post-construction wiring needed. + * *

Registers plugin settings defined in {@link ParquetSettings}. */ public class ParquetDataFormatPlugin extends Plugin implements DataFormatPlugin { @@ -86,14 +98,23 @@ public DataFormat getDataFormat() { } @Override - public IndexingExecutionEngine indexingEngine(IndexingEngineConfig engineSettings) { + public IndexingExecutionEngine indexingEngine(IndexingEngineConfig engineConfig, FormatChecksumStrategy checksumStrategy) { return new ParquetIndexingEngine( settings, dataFormat, - engineSettings.shardPath(), - () -> ArrowSchemaBuilder.getSchema(engineSettings.mapperService()), - engineSettings.indexSettings(), - threadPool + engineConfig.shardPath(), + () -> ArrowSchemaBuilder.getSchema(engineConfig.mapperService()), + engineConfig.indexSettings(), + threadPool, + checksumStrategy + ); + } + + @Override + public Map getFormatDescriptors(IndexSettings indexSettings) { + return Map.of( + ParquetDataFormat.PARQUET_DATA_FORMAT_NAME, + new DataFormatDescriptor(ParquetDataFormat.PARQUET_DATA_FORMAT_NAME, new PrecomputedChecksumStrategy()) ); } diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/ParquetFileMetadata.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/ParquetFileMetadata.java index 7ce133eac1126..d0baf747fd7e5 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/ParquetFileMetadata.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/ParquetFileMetadata.java @@ -12,12 +12,13 @@ * Metadata extracted from a Parquet file after the native writer is closed. * *

Returned by {@link RustBridge#finalizeWriter(String)} and {@link RustBridge#getFileMetadata(String)}. - * Contains the Parquet format version, total row count, and the creator identifier string - * embedded in the file footer. + * Contains the Parquet format version, total row count, the creator identifier string + * embedded in the file footer, and the whole-file CRC32 checksum computed during write. * * @param version Parquet format version number * @param numRows total number of rows written to the file * @param createdBy creator string from the Parquet file footer metadata + * @param crc32 whole-file CRC32 checksum (computed by the Rust writer during write) */ -public record ParquetFileMetadata(int version, long numRows, String createdBy) { +public record ParquetFileMetadata(int version, long numRows, String createdBy, long crc32) { } diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/RustBridge.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/RustBridge.java index d8cbf1b6c8b74..c9086cfe4e8e6 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/RustBridge.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/RustBridge.java @@ -55,6 +55,7 @@ public class RustBridge { ValueLayout.ADDRESS, ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, + ValueLayout.ADDRESS, ValueLayout.ADDRESS ) ); @@ -102,6 +103,7 @@ static ParquetFileMetadata finalizeWriter(String file) throws IOException { var f = call.str(file); var versionOut = call.intOut(); var numRowsOut = call.longOut(); + var crc32Out = call.longOut(); var out = call.outBuffer(1024); long rc = call.invokeIO( FINALIZE_WRITER, @@ -111,7 +113,8 @@ static ParquetFileMetadata finalizeWriter(String file) throws IOException { numRowsOut, out.data(), (long) out.capacity(), - out.lenOut() + out.lenOut(), + crc32Out ); if (rc == 1) return null; int createdByLen = out.actualLength(); @@ -120,7 +123,8 @@ static ParquetFileMetadata finalizeWriter(String file) throws IOException { numRowsOut.get(ValueLayout.JAVA_LONG, 0), createdByLen >= 0 ? new String(out.data().asSlice(0, createdByLen).toArray(ValueLayout.JAVA_BYTE), StandardCharsets.UTF_8) - : null + : null, + crc32Out.get(ValueLayout.JAVA_LONG, 0) ); } } @@ -145,7 +149,8 @@ public static ParquetFileMetadata getFileMetadata(String file) throws IOExceptio numRowsOut.get(ValueLayout.JAVA_LONG, 0), createdByLen >= 0 ? new String(out.data().asSlice(0, createdByLen).toArray(ValueLayout.JAVA_BYTE), StandardCharsets.UTF_8) - : null + : null, + 0L ); } } diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/engine/ParquetDataFormat.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/engine/ParquetDataFormat.java index 658ff6cb3b0a3..73cc70bde5ca9 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/engine/ParquetDataFormat.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/engine/ParquetDataFormat.java @@ -27,7 +27,8 @@ public class ParquetDataFormat extends DataFormat { /** Creates a new ParquetDataFormat. */ public ParquetDataFormat() {} - static final String PARQUET_DATA_FORMAT_NAME = "parquet"; + /** The parquet data format name constant. */ + public static final String PARQUET_DATA_FORMAT_NAME = "parquet"; @Override public String name() { diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/engine/ParquetIndexingEngine.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/engine/ParquetIndexingEngine.java index 437455767d1c1..a2e214c616c1a 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/engine/ParquetIndexingEngine.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/engine/ParquetIndexingEngine.java @@ -22,6 +22,8 @@ import org.opensearch.index.engine.exec.WriterFileSet; import org.opensearch.index.engine.exec.commit.IndexStoreProvider; import org.opensearch.index.shard.ShardPath; +import org.opensearch.index.store.FormatChecksumStrategy; +import org.opensearch.index.store.PrecomputedChecksumStrategy; import org.opensearch.parquet.bridge.RustBridge; import org.opensearch.parquet.memory.ArrowBufferPool; import org.opensearch.parquet.writer.ParquetDocumentInput; @@ -68,16 +70,17 @@ public class ParquetIndexingEngine implements IndexingExecutionEngine schemaSupplier, IndexSettings indexSettings, ThreadPool threadPool + ) { + this(settings, dataFormat, shardPath, schemaSupplier, indexSettings, threadPool, new PrecomputedChecksumStrategy()); + } + + /** + * Creates a new ParquetIndexingEngine with an externally provided checksum strategy. + * + *

Use this constructor when the checksum strategy is shared with the + * {@link org.opensearch.index.store.DataFormatAwareStoreDirectory} so that + * pre-computed CRC32 values registered during write are visible to the upload path. + * + * @param settings the node-level settings + * @param dataFormat the Parquet data format descriptor + * @param shardPath the shard path for file storage + * @param schemaSupplier supplier for the Arrow schema + * @param indexSettings the index-level settings + * @param threadPool the thread pool for background native writes + * @param checksumStrategy the checksum strategy to use (shared with the directory) + */ + public ParquetIndexingEngine( + Settings settings, + ParquetDataFormat dataFormat, + ShardPath shardPath, + Supplier schemaSupplier, + IndexSettings indexSettings, + ThreadPool threadPool, + FormatChecksumStrategy checksumStrategy ) { this.dataFormat = dataFormat; this.shardPath = shardPath; @@ -93,6 +123,16 @@ public ParquetIndexingEngine( this.bufferPool = new ArrowBufferPool(settings); this.settings = settings; this.threadPool = threadPool; + this.checksumStrategy = checksumStrategy; + } + + /** + * Returns the checksum strategy for this engine's Parquet files. + * Used by the upload path to retrieve pre-computed checksums. + */ + @Override + public FormatChecksumStrategy getChecksumStrategy() { + return checksumStrategy; } @Override @@ -102,7 +142,16 @@ public Writer createWriter(long writerGeneration) { dataFormat.name(), FILE_NAME_PREFIX + "_" + writerGeneration + FILE_NAME_EXT ); - return new ParquetWriter(filePath.toString(), writerGeneration, dataFormat, schemaSupplier.get(), bufferPool, settings, threadPool); + return new ParquetWriter( + filePath.toString(), + writerGeneration, + dataFormat, + schemaSupplier.get(), + bufferPool, + settings, + threadPool, + checksumStrategy + ); } @Override diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/writer/ParquetWriter.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/writer/ParquetWriter.java index f65628379fe50..283bdaf92f0c0 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/writer/ParquetWriter.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/writer/ParquetWriter.java @@ -14,6 +14,7 @@ import org.opensearch.index.engine.dataformat.WriteResult; import org.opensearch.index.engine.dataformat.Writer; import org.opensearch.index.engine.exec.WriterFileSet; +import org.opensearch.index.store.FormatChecksumStrategy; import org.opensearch.parquet.ParquetSettings; import org.opensearch.parquet.bridge.ParquetFileMetadata; import org.opensearch.parquet.engine.ParquetDataFormat; @@ -43,6 +44,7 @@ public class ParquetWriter implements Writer { private final long writerGeneration; private final ParquetDataFormat dataFormat; private final VSRManager vsrManager; + private final FormatChecksumStrategy checksumStrategy; /** * Creates a new ParquetWriter. @@ -54,6 +56,7 @@ public class ParquetWriter implements Writer { * @param bufferPool shared Arrow buffer pool * @param settings node settings for writer configuration * @param threadPool the thread pool for background native writes + * @param checksumStrategy strategy to register pre-computed checksums on */ public ParquetWriter( String file, @@ -62,12 +65,14 @@ public ParquetWriter( Schema schema, ArrowBufferPool bufferPool, Settings settings, - ThreadPool threadPool + ThreadPool threadPool, + FormatChecksumStrategy checksumStrategy ) { this.file = file; this.writerGeneration = writerGeneration; this.dataFormat = dataFormat; this.vsrManager = new VSRManager(file, schema, bufferPool, ParquetSettings.MAX_ROWS_PER_VSR.get(settings), threadPool); + this.checksumStrategy = checksumStrategy; } @Override @@ -83,10 +88,17 @@ public FileInfos flush() throws IOException { return FileInfos.empty(); } Path filePath = Path.of(file); + String fileName = filePath.getFileName().toString(); + + // Register the pre-computed CRC32 so the upload path can read it in O(1) + if (checksumStrategy != null && metadata.crc32() != 0) { + checksumStrategy.registerChecksum(fileName, metadata.crc32(), writerGeneration); + } + WriterFileSet writerFileSet = WriterFileSet.builder() .directory(filePath.getParent().getFileName()) .writerGeneration(writerGeneration) - .addFile(filePath.getFileName().toString()) + .addFile(fileName) .addNumRows(metadata.numRows()) .build(); return FileInfos.builder().putWriterFileSet(dataFormat, writerFileSet).build(); diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/Cargo.toml b/sandbox/plugins/parquet-data-format/src/main/rust/Cargo.toml index 1e867a6f76e3c..22466d27a3d60 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/Cargo.toml +++ b/sandbox/plugins/parquet-data-format/src/main/rust/Cargo.toml @@ -25,6 +25,7 @@ chrono = { workspace = true } mimalloc = { workspace = true } tempfile = { workspace = true } native-bridge-common = { workspace = true } +crc32fast = { workspace = true } [dev-dependencies] opensearch-parquet-format = { path = ".", features = ["test-utils"] } diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/ffm.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/ffm.rs index a1b2c4856e75a..f015a49110ec3 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/ffm.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/ffm.rs @@ -67,13 +67,15 @@ pub unsafe extern "C" fn parquet_finalize_writer( created_by_buf: *mut u8, created_by_buf_len: i64, created_by_len_out: *mut i64, + crc32_out: *mut i64, ) -> i64 { let filename = str_from_raw(file_ptr, file_len).map_err(|e| format!("parquet_finalize_writer: {}", e))?.to_string(); match NativeParquetWriter::finalize_writer(filename) { - Ok(Some(metadata)) => { - if !version_out.is_null() { *version_out = metadata.version; } - if !num_rows_out.is_null() { *num_rows_out = metadata.num_rows; } - if let Some(ref cb) = metadata.created_by { + Ok(Some(result)) => { + let fm = result.metadata.file_metadata(); + if !version_out.is_null() { *version_out = fm.version(); } + if !num_rows_out.is_null() { *num_rows_out = fm.num_rows(); } + if let Some(cb) = fm.created_by() { if !created_by_buf.is_null() && created_by_buf_len > 0 { let bytes = cb.as_bytes(); let n = bytes.len().min(created_by_buf_len as usize); @@ -83,6 +85,7 @@ pub unsafe extern "C" fn parquet_finalize_writer( } else if !created_by_len_out.is_null() { *created_by_len_out = -1; } + if !crc32_out.is_null() { *crc32_out = result.crc32 as i64; } Ok(0) } Ok(None) => Ok(1), diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/test_utils.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/test_utils.rs index b3840e81013ab..2a80157518ec8 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/test_utils.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/test_utils.rs @@ -100,7 +100,7 @@ pub fn create_mismatched_ffi_data() -> Result<(i64, i64), Box parquet::format::FileMetaData { +pub fn close_writer_and_get_metadata(filename: &str, schema_ptr: i64) -> crate::writer::FinalizeResult { let result = NativeParquetWriter::finalize_writer(filename.to_string()); cleanup_ffi_schema(schema_ptr); result.unwrap().unwrap() diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/tests/mod.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/tests/mod.rs index 913bf528c7834..9efcc961be225 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/tests/mod.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/tests/mod.rs @@ -9,6 +9,10 @@ use crate::test_utils::*; use crate::writer::{NativeParquetWriter, WRITER_MANAGER, FILE_MANAGER}; +use parquet::file::reader::FileReader; +use std::fs::File; +use std::io::Read; + #[test] fn test_create_writer_success() { let (_temp_dir, filename) = get_temp_file_path("test.parquet"); @@ -96,11 +100,11 @@ fn test_finalize_writer_success() { let (_schema, schema_ptr) = create_writer_and_assert_success(&filename); let result = NativeParquetWriter::finalize_writer(filename.clone()); assert!(result.is_ok()); - let metadata = result.unwrap(); - assert!(metadata.is_some()); - let metadata = metadata.unwrap(); - assert_eq!(metadata.num_rows, 0); - assert!(metadata.version > 0); + let finalize_result = result.unwrap(); + assert!(finalize_result.is_some()); + let finalize_result = finalize_result.unwrap(); + assert_eq!(finalize_result.metadata.file_metadata().num_rows(), 0); + assert!(finalize_result.metadata.file_metadata().version() > 0); assert!(!WRITER_MANAGER.contains_key(&filename)); assert!(FILE_MANAGER.contains_key(&filename)); FILE_MANAGER.remove(&filename); @@ -119,9 +123,10 @@ fn test_finalize_writer_with_data_returns_correct_metadata() { let result = NativeParquetWriter::finalize_writer(filename.clone()); assert!(result.is_ok()); let metadata = result.unwrap().unwrap(); - assert_eq!(metadata.num_rows, 6); - assert!(metadata.version > 0); - assert_eq!(metadata.schema.len(), 3); // root + 2 fields (id, name) + assert_eq!(metadata.metadata.file_metadata().num_rows(), 6); + assert!(metadata.metadata.file_metadata().version() > 0); + assert_eq!(metadata.metadata.file_metadata().schema_descr().num_columns(), 3); // root + 2 fields (id, name) + assert_ne!(metadata.crc32, 0, "CRC32 should be non-zero for a file with data"); FILE_MANAGER.remove(&filename); cleanup_ffi_schema(schema_ptr); } @@ -141,7 +146,7 @@ fn test_close_multiple_times_same_file() { assert!(result1.is_ok()); let metadata = result1.unwrap(); assert!(metadata.is_some()); - assert_eq!(metadata.unwrap().num_rows, 0); + assert_eq!(metadata.unwrap().metadata.num_rows, 0); assert!(!WRITER_MANAGER.contains_key(&filename)); assert!(FILE_MANAGER.contains_key(&filename)); let result2 = NativeParquetWriter::finalize_writer(filename.clone()); @@ -183,3 +188,122 @@ fn test_get_filtered_writer_memory_usage_with_writers() { close_writer_and_cleanup_schema(&filename1, schema_ptr1); close_writer_and_cleanup_schema(&filename2, schema_ptr2); } + + +/// Computes CRC32 of a file by reading it from disk in chunks. +/// This is the "re-read" baseline that the streaming checksum must match. +fn compute_file_crc32(path: &str) -> u32 { + let mut file = File::open(path).unwrap(); + let mut hasher = crc32fast::Hasher::new(); + let mut buf = [0u8; 64 * 1024]; + loop { + let n = file.read(&mut buf).unwrap(); + if n == 0 { + break; + } + hasher.update(&buf[..n]); + } + hasher.finalize() +} + +/// Verifies that the streaming CRC32 computed during write (via Crc32Writer) +/// exactly matches a CRC32 computed by re-reading the finalized file from disk. +/// +/// This proves the streaming approach is correct and eliminates the need for +/// a second I/O pass over the file. +#[test] +fn test_streaming_crc32_matches_reread_crc32_empty_file() { + let (_temp_dir, filename) = get_temp_file_path("crc32_empty.parquet"); + let (_schema, schema_ptr) = create_writer_and_assert_success(&filename); + + // Finalize with zero rows — still writes the Parquet magic bytes + footer + let result = NativeParquetWriter::finalize_writer(filename.clone()); + assert!(result.is_ok()); + let finalize_result = result.unwrap().unwrap(); + let streaming_crc32 = finalize_result.crc32; + + // Re-read the file and compute CRC32 independently + let reread_crc32 = compute_file_crc32(&filename); + + assert_eq!( + streaming_crc32, reread_crc32, + "Streaming CRC32 ({:#010x}) must match re-read CRC32 ({:#010x}) for empty Parquet file", + streaming_crc32, reread_crc32 + ); + assert_ne!(streaming_crc32, 0, "CRC32 should be non-zero even for an empty Parquet file (magic bytes + footer)"); + + FILE_MANAGER.remove(&filename); + cleanup_ffi_schema(schema_ptr); +} + +/// Verifies streaming CRC32 matches re-read CRC32 for a file with actual data. +/// Writes multiple batches to exercise the full write path (row groups, column +/// chunks, compression, bloom filters, footer). +#[test] +fn test_streaming_crc32_matches_reread_crc32_with_data() { + let (_temp_dir, filename) = get_temp_file_path("crc32_with_data.parquet"); + let (_schema, schema_ptr) = create_writer_and_assert_success(&filename); + + // Write 3 batches (9 rows total) to exercise multiple write() calls + for _ in 0..3 { + let (array_ptr, data_schema_ptr) = create_test_ffi_data().unwrap(); + NativeParquetWriter::write_data(filename.clone(), array_ptr, data_schema_ptr).unwrap(); + cleanup_ffi_data(array_ptr, data_schema_ptr); + } + + let result = NativeParquetWriter::finalize_writer(filename.clone()); + assert!(result.is_ok()); + let finalize_result = result.unwrap().unwrap(); + let streaming_crc32 = finalize_result.crc32; + + // Verify metadata is correct + assert_eq!(finalize_result.metadata.file_metadata().num_rows(), 9); + + // Re-read the file and compute CRC32 independently + let reread_crc32 = compute_file_crc32(&filename); + + assert_eq!( + streaming_crc32, reread_crc32, + "Streaming CRC32 ({:#010x}) must match re-read CRC32 ({:#010x}) for Parquet file with {} rows", + streaming_crc32, reread_crc32, finalize_result.metadata.file_metadata().num_rows() + ); + assert_ne!(streaming_crc32, 0, "CRC32 should be non-zero for a file with data"); + + // Verify the file is a valid Parquet file by reading it back + let file = File::open(&filename).unwrap(); + let reader = parquet::file::reader::SerializedFileReader::new(file).unwrap(); + assert_eq!(reader.metadata().file_metadata().num_rows(), 9); + + FILE_MANAGER.remove(&filename); + cleanup_ffi_schema(schema_ptr); +} + +/// Verifies that two different files produce different CRC32 values, +/// confirming the checksum is content-dependent and not a constant. +#[test] +fn test_streaming_crc32_differs_for_different_content() { + // File 1: empty + let (_temp_dir1, filename1) = get_temp_file_path("crc32_diff_a.parquet"); + let (_schema1, schema_ptr1) = create_writer_and_assert_success(&filename1); + let result1 = NativeParquetWriter::finalize_writer(filename1.clone()); + let crc32_empty = result1.unwrap().unwrap().crc32; + FILE_MANAGER.remove(&filename1); + cleanup_ffi_schema(schema_ptr1); + + // File 2: with data + let (_temp_dir2, filename2) = get_temp_file_path("crc32_diff_b.parquet"); + let (_schema2, schema_ptr2) = create_writer_and_assert_success(&filename2); + let (array_ptr, data_schema_ptr) = create_test_ffi_data().unwrap(); + NativeParquetWriter::write_data(filename2.clone(), array_ptr, data_schema_ptr).unwrap(); + cleanup_ffi_data(array_ptr, data_schema_ptr); + let result2 = NativeParquetWriter::finalize_writer(filename2.clone()); + let crc32_with_data = result2.unwrap().unwrap().crc32; + FILE_MANAGER.remove(&filename2); + cleanup_ffi_schema(schema_ptr2); + + assert_ne!( + crc32_empty, crc32_with_data, + "Empty file CRC32 ({:#010x}) should differ from file-with-data CRC32 ({:#010x})", + crc32_empty, crc32_with_data + ); +} diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/writer.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/writer.rs index ceb28a0fa0464..36bb2fe795d7d 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/writer.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/writer.rs @@ -14,15 +14,53 @@ use parquet::arrow::ArrowWriter; use parquet::basic::Compression; use parquet::file::properties::WriterProperties; use parquet::file::reader::{FileReader, SerializedFileReader}; -use parquet::file::metadata::ParquetMetaData; -use parquet::format::FileMetaData as FormatFileMetaData; use std::fs::File; +use std::io::Write; use std::sync::{Arc, Mutex}; -use crate::{log_info, log_error, log_debug}; +use crate::{log_error, log_debug}; + +/// A write wrapper that computes CRC32 as bytes flow through. +/// Wraps a File and tracks the running checksum without buffering. +pub struct Crc32Writer { + inner: File, + hasher: crc32fast::Hasher, +} + +impl Crc32Writer { + fn new(file: File) -> Self { + Self { + inner: file, + hasher: crc32fast::Hasher::new(), + } + } + + /// Finalizes and returns the CRC32 checksum of all bytes written. + fn checksum(&self) -> u32 { + self.hasher.clone().finalize() + } +} + +impl Write for Crc32Writer { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + let n = self.inner.write(buf)?; + self.hasher.update(&buf[..n]); + Ok(n) + } + + fn flush(&mut self) -> std::io::Result<()> { + self.inner.flush() + } +} + +/// Result from finalizing a writer: Parquet metadata + whole-file CRC32. +pub struct FinalizeResult { + pub metadata: parquet::file::metadata::ParquetMetaData, + pub crc32: u32, +} lazy_static! { - pub static ref WRITER_MANAGER: DashMap>>> = DashMap::new(); + pub static ref WRITER_MANAGER: DashMap>>> = DashMap::new(); pub static ref FILE_MANAGER: DashMap = DashMap::new(); } @@ -55,7 +93,8 @@ impl NativeParquetWriter { .set_bloom_filter_fpp(0.1) .set_bloom_filter_ndv(100000) .build(); - let writer = ArrowWriter::try_new(file, schema, Some(props))?; + let crc_writer = Crc32Writer::new(file); + let writer = ArrowWriter::try_new(crc_writer, schema, Some(props))?; WRITER_MANAGER.insert(filename, Arc::new(Mutex::new(writer))); Ok(()) } @@ -94,28 +133,19 @@ impl NativeParquetWriter { } } - pub fn finalize_writer(filename: String) -> Result, Box> { + pub fn finalize_writer(filename: String) -> Result, Box> { log_debug!("finalize_writer called for file: {}", filename); if let Some((_, writer_arc)) = WRITER_MANAGER.remove(&filename) { match Arc::try_unwrap(writer_arc) { Ok(mutex) => { - let writer = mutex.into_inner().unwrap(); - let parquet_metadata = writer.close()?; + let mut writer = mutex.into_inner().unwrap(); + let parquet_metadata = writer.finish()?; let file_metadata = parquet_metadata.file_metadata(); - log_debug!("Successfully closed writer for file: {}, num_rows={}", filename, file_metadata.num_rows()); - let format_metadata = FormatFileMetaData { - version: file_metadata.version(), - num_rows: file_metadata.num_rows(), - created_by: file_metadata.created_by().map(|s| s.to_string()), - schema: vec![], - row_groups: vec![], - key_value_metadata: None, - encryption_algorithm: None, - footer_signing_key_metadata: None, - column_orders: None, - }; - Ok(Some(format_metadata)) + log_debug!("Successfully finalized writer for file: {}, num_rows={}", filename, file_metadata.num_rows()); + let crc32 = writer.inner().checksum(); + log_debug!("CRC32 for file {}: {:#010x}", filename, crc32); + Ok(Some(FinalizeResult { metadata: parquet_metadata, crc32 })) } Err(_) => { log_error!("ERROR: Writer still in use for file: {}", filename); diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/tests/writer_integration_tests.rs b/sandbox/plugins/parquet-data-format/src/main/rust/tests/writer_integration_tests.rs index f7a06e87491ea..8a0bc1c6c8778 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/tests/writer_integration_tests.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/tests/writer_integration_tests.rs @@ -25,17 +25,17 @@ fn test_complete_writer_lifecycle() { } let metadata = close_writer_and_get_metadata(&filename, schema_ptr); - assert_eq!(metadata.num_rows, 9); // 3 batches × 3 rows - assert!(metadata.version > 0); - assert_eq!(metadata.schema.len(), 3); // root + 2 fields + assert_eq!(metadata.metadata.file_metadata().num_rows(), 9); // 3 batches × 3 rows + assert!(metadata.metadata.file_metadata().version() > 0); + assert_eq!(metadata.metadata.file_metadata().schema_descr().num_columns(), 3); // root + 2 fields assert!(NativeParquetWriter::sync_to_disk(filename.clone()).is_ok()); assert!(file_path.exists()); assert!(file_path.metadata().unwrap().len() > 0); let read_metadata = NativeParquetWriter::get_file_metadata(filename.clone()).unwrap(); - assert_eq!(read_metadata.num_rows(), metadata.num_rows); - assert_eq!(read_metadata.version(), metadata.version); + assert_eq!(read_metadata.num_rows(), metadata.metadata.file_metadata().num_rows()); + assert_eq!(read_metadata.version(), metadata.metadata.file_metadata().version()); } #[test] @@ -181,7 +181,7 @@ fn test_concurrent_complete_writer_lifecycle() { if write_ok { if let Ok(Some(metadata)) = NativeParquetWriter::finalize_writer(filename.clone()) { - if metadata.num_rows == 3 + if metadata.metadata.file_metadata().num_rows() == 3 && NativeParquetWriter::sync_to_disk(filename.clone()).is_ok() && file_path.exists() { diff --git a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/writer/ParquetWriterTests.java b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/writer/ParquetWriterTests.java index 076e6f71e34e2..7fa90cf358ed5 100644 --- a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/writer/ParquetWriterTests.java +++ b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/writer/ParquetWriterTests.java @@ -71,7 +71,16 @@ public void tearDown() throws Exception { public void testAddDocReturnsSuccess() throws Exception { String filePath = createTempDir().resolve("success.parquet").toString(); - ParquetWriter writer = new ParquetWriter(filePath, 1L, new ParquetDataFormat(), schema, bufferPool, Settings.EMPTY, threadPool); + ParquetWriter writer = new ParquetWriter( + filePath, + 1L, + new ParquetDataFormat(), + schema, + bufferPool, + Settings.EMPTY, + threadPool, + null + ); ParquetDocumentInput doc = new ParquetDocumentInput(); doc.addField(idField, 1); @@ -85,7 +94,16 @@ public void testAddDocReturnsSuccess() throws Exception { public void testSingleDocumentFlush() throws Exception { String filePath = createTempDir().resolve("single.parquet").toString(); - ParquetWriter writer = new ParquetWriter(filePath, 1L, new ParquetDataFormat(), schema, bufferPool, Settings.EMPTY, threadPool); + ParquetWriter writer = new ParquetWriter( + filePath, + 1L, + new ParquetDataFormat(), + schema, + bufferPool, + Settings.EMPTY, + threadPool, + null + ); ParquetDocumentInput doc = new ParquetDocumentInput(); doc.addField(idField, 42); @@ -100,7 +118,16 @@ public void testSingleDocumentFlush() throws Exception { public void testMultipleDocumentsFlush() throws Exception { String filePath = createTempDir().resolve("multi.parquet").toString(); - ParquetWriter writer = new ParquetWriter(filePath, 1L, new ParquetDataFormat(), schema, bufferPool, Settings.EMPTY, threadPool); + ParquetWriter writer = new ParquetWriter( + filePath, + 1L, + new ParquetDataFormat(), + schema, + bufferPool, + Settings.EMPTY, + threadPool, + null + ); for (int i = 0; i < 10; i++) { ParquetDocumentInput doc = new ParquetDocumentInput(); @@ -119,13 +146,31 @@ public void testMultipleDocumentsFlush() throws Exception { public void testFlushWithNoDocuments() throws Exception { String filePath = createTempDir().resolve("empty.parquet").toString(); - ParquetWriter writer = new ParquetWriter(filePath, 1L, new ParquetDataFormat(), schema, bufferPool, Settings.EMPTY, threadPool); + ParquetWriter writer = new ParquetWriter( + filePath, + 1L, + new ParquetDataFormat(), + schema, + bufferPool, + Settings.EMPTY, + threadPool, + null + ); assertEquals(FileInfos.empty(), writer.flush()); } public void testSyncAfterFlush() throws Exception { String filePath = createTempDir().resolve("sync.parquet").toString(); - ParquetWriter writer = new ParquetWriter(filePath, 1L, new ParquetDataFormat(), schema, bufferPool, Settings.EMPTY, threadPool); + ParquetWriter writer = new ParquetWriter( + filePath, + 1L, + new ParquetDataFormat(), + schema, + bufferPool, + Settings.EMPTY, + threadPool, + null + ); ParquetDocumentInput doc = new ParquetDocumentInput(); doc.addField(idField, 1); diff --git a/server/src/main/java/org/opensearch/index/IndexModule.java b/server/src/main/java/org/opensearch/index/IndexModule.java index 589a98b6a50d2..28f55bd6f8d14 100644 --- a/server/src/main/java/org/opensearch/index/IndexModule.java +++ b/server/src/main/java/org/opensearch/index/IndexModule.java @@ -84,6 +84,7 @@ import org.opensearch.index.shard.SearchOperationListener; import org.opensearch.index.shard.ShardPath; import org.opensearch.index.similarity.SimilarityService; +import org.opensearch.index.store.DataFormatAwareStoreDirectoryFactory; import org.opensearch.index.store.DefaultCompositeDirectoryFactory; import org.opensearch.index.store.FsDirectoryFactory; import org.opensearch.index.store.Store; @@ -274,6 +275,7 @@ public final class IndexModule { private final Map> similarities = new HashMap<>(); private final Map directoryFactories; private final Map compositeDirectoryFactories; + private final Map dataFormatAwareStoreDirectoryFactories; private final SetOnce> forceQueryCacheProvider = new SetOnce<>(); private final List searchOperationListeners = new ArrayList<>(); private final List indexOperationListeners = new ArrayList<>(); @@ -307,7 +309,8 @@ public IndexModule( final Map recoveryStateFactories, final Map storeFactories, final FileCache fileCache, - final CompositeIndexSettings compositeIndexSettings + final CompositeIndexSettings compositeIndexSettings, + final Map dataFormatAwareStoreDirectoryFactories ) { this.indexSettings = indexSettings; this.analysisRegistry = analysisRegistry; @@ -317,6 +320,7 @@ public IndexModule( this.indexOperationListeners.add(new IndexingSlowLog(indexSettings)); this.directoryFactories = Collections.unmodifiableMap(directoryFactories); this.compositeDirectoryFactories = Collections.unmodifiableMap(compositeDirectoryFactories); + this.dataFormatAwareStoreDirectoryFactories = Collections.unmodifiableMap(dataFormatAwareStoreDirectoryFactories); this.allowExpensiveQueries = allowExpensiveQueries; this.expressionResolver = expressionResolver; this.recoveryStateFactories = recoveryStateFactories; @@ -348,7 +352,8 @@ public IndexModule( recoveryStateFactories, Collections.emptyMap(), null, - null + null, + Collections.emptyMap() ); } @@ -946,6 +951,10 @@ public IndexService newIndexService( indexSettings, compositeDirectoryFactories ); + final DataFormatAwareStoreDirectoryFactory dataFormatAwareStoreDirectoryFactory = getDataFormatAwareStoreDirectoryFactory( + indexSettings, + dataFormatAwareStoreDirectoryFactories + ); final IndexStorePlugin.RecoveryStateFactory recoveryStateFactory = getRecoveryStateFactory(indexSettings, recoveryStateFactories); QueryCache queryCache = null; IndexAnalyzers indexAnalyzers = null; @@ -1009,7 +1018,8 @@ public IndexService newIndexService( segmentReplicationStatsProvider, clusterDefaultMaxMergeAtOnceSupplier, clusterMergeSchedulerConfig, - dataFormatRegistry + dataFormatRegistry, + dataFormatAwareStoreDirectoryFactory ); success = true; return indexService; @@ -1079,6 +1089,16 @@ private static IndexStorePlugin.CompositeDirectoryFactory getCompositeDirectoryF return factory; } + private static DataFormatAwareStoreDirectoryFactory getDataFormatAwareStoreDirectoryFactory( + final IndexSettings indexSettings, + final Map dataFormatAwareStoreDirectoryFactories + ) { + if (dataFormatAwareStoreDirectoryFactories.isEmpty()) { + return null; + } + return dataFormatAwareStoreDirectoryFactories.get("default"); + } + private static IndexStorePlugin.RecoveryStateFactory getRecoveryStateFactory( final IndexSettings indexSettings, final Map recoveryStateFactories diff --git a/server/src/main/java/org/opensearch/index/IndexService.java b/server/src/main/java/org/opensearch/index/IndexService.java index ede50d4bb96e7..285a3e9459921 100644 --- a/server/src/main/java/org/opensearch/index/IndexService.java +++ b/server/src/main/java/org/opensearch/index/IndexService.java @@ -96,6 +96,8 @@ import org.opensearch.index.shard.ShardNotInPrimaryModeException; import org.opensearch.index.shard.ShardPath; import org.opensearch.index.similarity.SimilarityService; +import org.opensearch.index.store.DataFormatAwareStoreDirectory; +import org.opensearch.index.store.DataFormatAwareStoreDirectoryFactory; import org.opensearch.index.store.RemoteSegmentStoreDirectoryFactory; import org.opensearch.index.store.Store; import org.opensearch.index.store.remote.filecache.FileCache; @@ -159,6 +161,7 @@ public class IndexService extends AbstractIndexComponent implements IndicesClust private final ShardStoreDeleter shardStoreDeleter; private final IndexStorePlugin.DirectoryFactory directoryFactory; private final IndexStorePlugin.CompositeDirectoryFactory compositeDirectoryFactory; + private final DataFormatAwareStoreDirectoryFactory dataFormatAwareStoreDirectoryFactory; private final IndexStorePlugin.DirectoryFactory remoteDirectoryFactory; private final IndexStorePlugin.RecoveryStateFactory recoveryStateFactory; private final CheckedFunction readerWrapper; @@ -258,7 +261,8 @@ public IndexService( Function segmentReplicationStatsProvider, Supplier clusterDefaultMaxMergeAtOnceSupplier, ClusterMergeSchedulerConfig clusterMergeSchedulerConfig, - DataFormatRegistry dataFormatRegistry + DataFormatRegistry dataFormatRegistry, + DataFormatAwareStoreDirectoryFactory dataFormatAwareStoreDirectoryFactory ) { super(indexSettings); this.storeFactory = storeFactory; @@ -327,6 +331,7 @@ public IndexService( this.nodeEnv = nodeEnv; this.directoryFactory = directoryFactory; this.compositeDirectoryFactory = compositeDirectoryFactory; + this.dataFormatAwareStoreDirectoryFactory = dataFormatAwareStoreDirectoryFactory; this.remoteDirectoryFactory = remoteDirectoryFactory; this.recoveryStateFactory = recoveryStateFactory; this.engineFactory = Objects.requireNonNull(engineFactory); @@ -461,6 +466,7 @@ public IndexService( (shardId) -> ReplicationStats.empty(), clusterDefaultMaxMergeAtOnce, clusterMergeSchedulerConfig, + null, null ); } @@ -769,8 +775,11 @@ protected void closeInternal() { fileCache, threadPool ); - } else { + } else if (!this.indexSettings.isPluggableDataFormatEnabled()) { directory = directoryFactory.newDirectory(this.indexSettings, path); + } else { + // Will be enabled in case of formatAware indices. + directory = createDataFormatAwareStoreDirectory(shardId, path); } store = storeFactory.newStore( shardId, @@ -1322,6 +1331,26 @@ public boolean isForceExecution() { rescheduleRefreshTasks(); } + /** + * Creates DataFormatAwareStoreDirectory using the factory if available, otherwise fallback to Store's internal creation. + * This method centralizes the directory creation logic and enables plugin-based format discovery. + */ + private DataFormatAwareStoreDirectory createDataFormatAwareStoreDirectory(ShardId shardId, ShardPath shardPath) throws IOException { + if (dataFormatAwareStoreDirectoryFactory != null) { + logger.debug("Using DataFormatAwareStoreDirectoryFactory to create directory for shard path: {}", shardPath); + return dataFormatAwareStoreDirectoryFactory.newDataFormatAwareStoreDirectory( + indexSettings, + shardId, + shardPath, + directoryFactory, + dataFormatRegistry + ); + } + + logger.debug("No DataFormatAwareStoreDirectoryFactory available, Store will handle internal creation for: {}", shardPath); + return null; + } + private void updateFsyncTaskIfNecessary() { if (indexSettings.getTranslogDurability() == Translog.Durability.REQUEST) { try { 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 eb03ff54c0e11..bd55931f9763b 100644 --- a/server/src/main/java/org/opensearch/index/engine/EngineBackedIndexer.java +++ b/server/src/main/java/org/opensearch/index/engine/EngineBackedIndexer.java @@ -9,6 +9,7 @@ package org.opensearch.index.engine; import org.apache.lucene.index.IndexCommit; +import org.apache.lucene.index.SegmentInfos; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.common.concurrent.GatedCloseable; import org.opensearch.common.unit.TimeValue; @@ -16,6 +17,7 @@ import org.opensearch.index.VersionType; import org.opensearch.index.engine.exec.Indexer; import org.opensearch.index.engine.exec.coord.CatalogSnapshot; +import org.opensearch.index.engine.exec.coord.SegmentInfosCatalogSnapshot; import org.opensearch.index.mapper.DocumentMapperForType; import org.opensearch.index.mapper.SourceToParse; import org.opensearch.index.merge.MergeStats; @@ -382,9 +384,9 @@ public long getNativeBytesUsed() { @ExperimentalApi @Override public GatedCloseable acquireSnapshot() { - // TODO: Replace with a SegmentInfosCatalogSnapshot - // For now we throw an exception as this is not yet implemented - throw new UnsupportedOperationException("acquireSnapshot is not supported in EngineBackedIndexer"); + GatedCloseable segmentInfosRef = engine.getSegmentInfosSnapshot(); + SegmentInfosCatalogSnapshot snapshot = new SegmentInfosCatalogSnapshot(segmentInfosRef.get()); + return new GatedCloseable<>(snapshot, segmentInfosRef::close); } @Override diff --git a/server/src/main/java/org/opensearch/index/engine/dataformat/DataFormatDescriptor.java b/server/src/main/java/org/opensearch/index/engine/dataformat/DataFormatDescriptor.java new file mode 100644 index 0000000000000..0df1498a23b41 --- /dev/null +++ b/server/src/main/java/org/opensearch/index/engine/dataformat/DataFormatDescriptor.java @@ -0,0 +1,60 @@ +/* + * 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.dataformat; + +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.index.store.FormatChecksumStrategy; + +/** + * Describes the static capabilities of a data format, including its default checksum + * strategy and format name. Provided by {@link DataFormatPlugin} implementations and + * consumed by DataFormatAwareStoreDirectory and DataFormatAwareRemoteDirectory. + * + *

The checksum strategy here is the default fallback — a full-file scan. + * At runtime, the {@link IndexingExecutionEngine} may override this with a more + * efficient strategy (e.g., {@link org.opensearch.index.store.PrecomputedChecksumStrategy}) + * via {@link org.opensearch.index.store.DataFormatAwareStoreDirectory#registerChecksumStrategy}. + * + * @opensearch.experimental + */ +@ExperimentalApi +public class DataFormatDescriptor { + + private final String formatName; + private final FormatChecksumStrategy checksumStrategy; + + /** + * Creates a new DataFormatDescriptor. + * + * @param formatName the format name (e.g., "parquet") + * @param checksumStrategy the default checksum strategy for this format + */ + public DataFormatDescriptor(String formatName, FormatChecksumStrategy checksumStrategy) { + this.formatName = formatName; + this.checksumStrategy = checksumStrategy; + } + + /** + * Returns the format name. + * + * @return the format name + */ + public String getFormatName() { + return formatName; + } + + /** + * Returns the default checksum strategy for this format. + * + * @return the checksum strategy + */ + public FormatChecksumStrategy getChecksumStrategy() { + return checksumStrategy; + } +} diff --git a/server/src/main/java/org/opensearch/index/engine/dataformat/DataFormatPlugin.java b/server/src/main/java/org/opensearch/index/engine/dataformat/DataFormatPlugin.java index 2db9891a8efcb..90228de27d662 100644 --- a/server/src/main/java/org/opensearch/index/engine/dataformat/DataFormatPlugin.java +++ b/server/src/main/java/org/opensearch/index/engine/dataformat/DataFormatPlugin.java @@ -9,6 +9,10 @@ package org.opensearch.index.engine.dataformat; import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.index.IndexSettings; +import org.opensearch.index.store.FormatChecksumStrategy; + +import java.util.Map; /** * Plugin interface for providing custom data format implementations. @@ -30,8 +34,24 @@ public interface DataFormatPlugin { /** * Creates the indexing engine for the data format. This should be instantiated per shard. * - * @param settings the engine initialization settings + * @param settings the engine initialization settings + * @param checksumStrategy the checksum strategy owned by the directory for this format, + * or null if not available. Engines that pre-compute checksums + * during write should register into this instance so the upload + * path can retrieve them in O(1). * @return the indexing execution engine instance */ - IndexingExecutionEngine indexingEngine(IndexingEngineConfig settings); + IndexingExecutionEngine indexingEngine(IndexingEngineConfig settings, FormatChecksumStrategy checksumStrategy); + + /** + * Returns format descriptors for this plugin, filtered by the given index settings. + * Each entry maps a format name to its {@link DataFormatDescriptor} containing the + * default checksum strategy and format name. + * + * @param indexSettings the index settings used to determine active formats + * @return map of format name to descriptor + */ + default Map getFormatDescriptors(IndexSettings indexSettings) { + return Map.of(); + } } diff --git a/server/src/main/java/org/opensearch/index/engine/dataformat/DataFormatRegistry.java b/server/src/main/java/org/opensearch/index/engine/dataformat/DataFormatRegistry.java index 5d6e5b7d2a146..1db3323b80790 100644 --- a/server/src/main/java/org/opensearch/index/engine/dataformat/DataFormatRegistry.java +++ b/server/src/main/java/org/opensearch/index/engine/dataformat/DataFormatRegistry.java @@ -15,6 +15,7 @@ import org.opensearch.index.engine.exec.commit.IndexStoreProvider; import org.opensearch.index.mapper.MapperService; import org.opensearch.index.shard.ShardPath; +import org.opensearch.index.store.FormatChecksumStrategy; import org.opensearch.plugins.PluginsService; import org.opensearch.plugins.SearchBackEndPlugin; @@ -36,6 +37,9 @@ @ExperimentalApi public class DataFormatRegistry { + /** Index setting name that specifies the active pluggable data format. */ + public static final String PLUGGABLE_DATAFORMAT_SETTING = "pluggable_dataformat"; + /** Map from data format to the plugin that provides its indexing engine. */ private final Map dataFormatPluginRegistry; @@ -100,7 +104,10 @@ public DataFormatRegistry(PluginsService pluginsService) { if (plugin == null) { throw new IllegalArgumentException("No plugin registered for DataFormat [" + format.name() + "]"); } - return plugin.indexingEngine(settings); + Map descriptors = plugin.getFormatDescriptors(settings.indexSettings()); + DataFormatDescriptor descriptor = descriptors.get(format.name()); + FormatChecksumStrategy checksumStrategy = descriptor != null ? descriptor.getChecksumStrategy() : null; + return plugin.indexingEngine(settings, checksumStrategy); } public DataFormat format(String name) { @@ -133,12 +140,34 @@ public List supportsCapability(String fieldType, FieldTypeCapabiliti /** * Returns an unmodifiable view of all registered data formats and their plugins. * - * @return unmodifiable map of data formats to plugins + * @return unmodifiable set of data formats */ public Set getRegisteredFormats() { return Set.copyOf(dataFormatPluginRegistry.keySet()); } + /** + * Returns format descriptors for the active data format of the given index. + * Resolves the data format from index settings via the {@code pluggable_dataformat} setting, + * then delegates to {@link DataFormatPlugin#getFormatDescriptors(IndexSettings)}. + * + * @param indexSettings the index settings used to determine the active data format + * @return unmodifiable map of format name to descriptor, or empty map if no pluggable data format is configured + */ + public Map getFormatDescriptors(IndexSettings indexSettings) { + String dataformatName = indexSettings.getSettings().get(PLUGGABLE_DATAFORMAT_SETTING); + if (dataformatName != null) { + DataFormat format = dataFormats.get(dataformatName); + if (format != null) { + DataFormatPlugin plugin = dataFormatPluginRegistry.get(format); + if (plugin != null) { + return plugin.getFormatDescriptors(indexSettings); + } + } + } + return Map.of(); + } + /** * Creates {@link EngineReaderManager} instances for all applicable data formats based on index settings/mappings. * Each reader manager is instantiated by applying the store provider and shard path to the factory registered diff --git a/server/src/main/java/org/opensearch/index/engine/dataformat/IndexingExecutionEngine.java b/server/src/main/java/org/opensearch/index/engine/dataformat/IndexingExecutionEngine.java index 8f8b7c6a414fd..f64bdea610810 100644 --- a/server/src/main/java/org/opensearch/index/engine/dataformat/IndexingExecutionEngine.java +++ b/server/src/main/java/org/opensearch/index/engine/dataformat/IndexingExecutionEngine.java @@ -10,6 +10,7 @@ import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.index.engine.exec.commit.IndexStoreProvider; +import org.opensearch.index.store.FormatChecksumStrategy; import java.io.IOException; import java.util.Collection; @@ -99,4 +100,19 @@ default long getNativeBytesUsed() { * @return the store provider, or null if this engine does not expose one */ IndexStoreProvider getProvider(); + + /** + * Returns the checksum strategy used by this engine, if any. + * + *

Engines that pre-compute checksums during write (e.g., Parquet computing CRC32 + * in the native writer) return their strategy here so it can be wired into the + * {@link org.opensearch.index.store.DataFormatAwareStoreDirectory} at shard init time. + * This allows the upload path to retrieve pre-computed checksums in O(1) instead of + * re-reading the entire file. + * + * @return the checksum strategy, or {@code null} if this engine does not pre-compute checksums + */ + default FormatChecksumStrategy getChecksumStrategy() { + return null; + } } diff --git a/server/src/main/java/org/opensearch/index/engine/exec/FileMetadata.java b/server/src/main/java/org/opensearch/index/engine/exec/FileMetadata.java deleted file mode 100644 index 71b85e0c2a4c6..0000000000000 --- a/server/src/main/java/org/opensearch/index/engine/exec/FileMetadata.java +++ /dev/null @@ -1,106 +0,0 @@ -/* - * 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.util.Objects; - -/** - * Represents metadata for a file in the index, including its data format and filename. - * Files can be in different formats (e.g., "lucene", "metadata") and this class provides - * a unified way to represent and serialize file information across the system. - * - * @opensearch.experimental - */ -@ExperimentalApi -public class FileMetadata { - - /** - * Delimiter used to separate filename and data format in serialized form. - */ - public static final String DELIMITER = ":::"; - private static final String METADATA_KEY = "metadata"; - - private final String file; - private final String dataFormat; - - /** - * Constructs a FileMetadata with explicit data format and filename. - * - * @param dataFormat the data format identifier (e.g., "lucene", "metadata") - * @param file the filename - */ - public FileMetadata(String dataFormat, String file) { - this.file = file; - this.dataFormat = dataFormat; - } - - /** - * Constructs a FileMetadata by parsing a serialized data-format-aware filename. - * The format is "filename:::dataFormat". If no delimiter is present and the filename - * starts with "metadata", it's treated as a metadata file. Otherwise, defaults to "lucene". - * - * @param dataFormatAwareFile the serialized filename with optional data format - */ - public FileMetadata(String dataFormatAwareFile) { - if (!dataFormatAwareFile.contains(DELIMITER) && dataFormatAwareFile.startsWith(METADATA_KEY)) { - this.dataFormat = "metadata"; - this.file = dataFormatAwareFile; - return; - } - String[] parts = dataFormatAwareFile.split(DELIMITER); - this.dataFormat = (parts.length == 1) ? "lucene" : parts[1]; - this.file = parts[0]; - } - - /** - * Serializes this FileMetadata to a string in the format "filename:::dataFormat". - * - * @return the serialized representation - */ - public String serialize() { - return file + DELIMITER + dataFormat; - } - - @Override - public String toString() { - return serialize(); - } - - /** - * Returns the filename. - * - * @return the filename - */ - public String file() { - return file; - } - - /** - * Returns the data format identifier. - * - * @return the data format (e.g., "lucene", "metadata") - */ - public String dataFormat() { - return dataFormat; - } - - @Override - public boolean equals(Object o) { - if (o == null || getClass() != o.getClass()) return false; - FileMetadata that = (FileMetadata) o; - return Objects.equals(file, that.file) && Objects.equals(dataFormat, that.dataFormat); - } - - @Override - public int hashCode() { - return Objects.hash(file, dataFormat); - } -} diff --git a/server/src/main/java/org/opensearch/index/engine/exec/Segment.java b/server/src/main/java/org/opensearch/index/engine/exec/Segment.java index 576d871832dde..251e03968f02a 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/Segment.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/Segment.java @@ -22,6 +22,7 @@ /** * Represents a segment in the catalog snapshot containing files grouped by data format. * Each segment has a unique generation number and maintains searchable files organized by their data format type. + * This class is serializable and can be transmitted across nodes for replication and recovery operations. */ @ExperimentalApi public record Segment(long generation, Map dfGroupedSearchableFiles) implements Writeable { 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 ca7169cc535fd..62fe6838c202f 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 @@ -191,8 +191,9 @@ public CatalogSnapshot cloneNoAcquire() { * Sets user-defined metadata for this catalog snapshot. * * @param userData map of user data key-value pairs + * @param commitData whether this is commit-level user data */ - public abstract void setUserData(Map userData); + public abstract void setUserData(Map userData, boolean commitData); /** * Creates a deep copy of this catalog snapshot. The cloned snapshot starts with a fresh reference count of 1. @@ -201,4 +202,42 @@ public CatalogSnapshot cloneNoAcquire() { * @return a new {@link CatalogSnapshot} with the same logical state */ public abstract CatalogSnapshot clone(); + + /** + * Returns the major version of the format that wrote the given file. + * For Lucene files, this is the Lucene major version from SegmentInfo. + * For non-Lucene files (e.g., parquet), this is the format-specific version. + * + * @param file the file name + * @return the format major version + */ + public abstract int getFormatVersionForFile(String file); + + /** + * Serializes this CatalogSnapshot into SegmentInfos bytes for the remote metadata file. + * Each subclass knows its own serialization format: + * + * TODO: When CompositeEngineCatalogSnapshot is added, implement this method + * creating synthetic SegmentInfos with CatalogSnapshot serialized into userData. + * + * @return serialized bytes + * @throws IOException in case of I/O error + */ + public abstract byte[] serialize() throws IOException; + + /** + * Returns the canonical file names for upload to remote store. + * Each subclass formats names appropriately for its data format: + *

    + *
  • {@link SegmentInfosCatalogSnapshot}: plain Lucene file names (e.g., {@code "_0.cfe"})
    • + *
  • {@link DataformatAwareCatalogSnapshot}: serialized format-aware names + * (e.g., {@code "parquet/data.parquet"}) for non-lucene files
    • + *
+ * + * @param includeSegmentsFile whether to include the segments file in the returned collection + * @return collection of file name strings ready for upload + * @throws IOException in case of I/O error + */ + public abstract Collection getFiles(boolean includeSegmentsFile) throws IOException; + } diff --git a/server/src/main/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshot.java b/server/src/main/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshot.java index 86348e82099cf..cb243e952731c 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshot.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshot.java @@ -16,6 +16,7 @@ import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.index.engine.exec.Segment; import org.opensearch.index.engine.exec.WriterFileSet; +import org.opensearch.index.store.FileMetadata; import java.io.IOException; import java.util.ArrayList; @@ -133,7 +134,7 @@ public Map getUserData() { } @Override - public void setUserData(Map userData) { + public void setUserData(Map userData, boolean commitData) { this.userData = Map.copyOf(userData); } @@ -187,6 +188,20 @@ public DataformatAwareCatalogSnapshot clone() { return new DataformatAwareCatalogSnapshot(id, generation, version, segments, lastWriterGeneration, userData); } + @Override + public int getFormatVersionForFile(String file) { + // TODO: Return the actual format-specific version the file was written with. + // For lucene files, this should come from a per-segment version map populated + // by the composite engine (which has access to SegmentInfos). For non-lucene + // files, each DataFormat should provide its own version. + return org.opensearch.Version.CURRENT.major; + } + + @Override + public byte[] serialize() throws IOException { + throw new UnsupportedOperationException("DataformatAwareCatalogSnapshot does not support serialize()"); + } + @Override protected void closeInternal() { closed.set(true); @@ -199,4 +214,18 @@ protected void closeInternal() { public boolean isClosed() { return closed.get(); } + + @Override + public Collection getFiles(boolean includeSegmentsFile) throws IOException { + List fileNames = new ArrayList<>(); + for (Segment segment : segments) { + for (Map.Entry entry : segment.dfGroupedSearchableFiles().entrySet()) { + String formatName = entry.getKey(); + for (String file : entry.getValue().files()) { + fileNames.add(FileMetadata.serialize(formatName, file)); + } + } + } + return fileNames; + } } diff --git a/server/src/main/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshot.java b/server/src/main/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshot.java index 7b43e9d93f616..6e2cad0c13d43 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshot.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshot.java @@ -8,6 +8,8 @@ package org.opensearch.index.engine.exec.coord; +import org.apache.lucene.index.SegmentCommitInfo; +import org.apache.lucene.index.SegmentInfo; import org.apache.lucene.index.SegmentInfos; import org.apache.lucene.store.BufferedChecksumIndexInput; import org.apache.lucene.store.ByteBuffersDataOutput; @@ -18,9 +20,11 @@ import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.index.engine.exec.Segment; import org.opensearch.index.engine.exec.WriterFileSet; +import org.opensearch.index.remote.RemoteStoreUtils; import java.io.IOException; import java.util.Collection; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Set; @@ -40,6 +44,7 @@ public class SegmentInfosCatalogSnapshot extends CatalogSnapshot { private static final String CATALOG_SNAPSHOT_KEY = "_segment_infos_catalog_snapshot_"; private final SegmentInfos segmentInfos; + private final Map segmentFileVersionMap; /** * Constructs a new SegmentInfosCatalogSnapshot wrapping the given SegmentInfos. @@ -49,6 +54,7 @@ public class SegmentInfosCatalogSnapshot extends CatalogSnapshot { public SegmentInfosCatalogSnapshot(SegmentInfos segmentInfos) { super(CATALOG_SNAPSHOT_KEY + segmentInfos.getGeneration(), segmentInfos.getGeneration(), segmentInfos.getVersion()); this.segmentInfos = segmentInfos; + this.segmentFileVersionMap = buildSegmentToLuceneVersionMap(); } /** @@ -66,6 +72,7 @@ public SegmentInfosCatalogSnapshot(StreamInput in) throws IOException { new BufferedChecksumIndexInput(new ByteArrayIndexInput("SegmentInfos", segmentInfosBytes)), 0L ); + this.segmentFileVersionMap = buildSegmentToLuceneVersionMap(); } /** @@ -123,8 +130,8 @@ public void writeTo(StreamOutput out) throws IOException { } @Override - public void setUserData(Map userData) { - // No-op for SegmentInfosCatalogSnapshot + public void setUserData(Map userData, boolean commitData) { + segmentInfos.setUserData(userData, commitData); } @Override @@ -134,6 +141,60 @@ protected void closeInternal() { @Override public SegmentInfosCatalogSnapshot clone() { - return new SegmentInfosCatalogSnapshot(segmentInfos); + return new SegmentInfosCatalogSnapshot(segmentInfos.clone()); + } + + @Override + public CatalogSnapshot cloneNoAcquire() { + return new SegmentInfosCatalogSnapshot(segmentInfos.clone()); + } + + /** + * Returns the Lucene major version that wrote the given segment file by looking it up + * from the SegmentInfos. Falls back to the SegmentInfos commit version for the segments + * file itself, or to the .si file's version for other unmapped files. + */ + @Override + public int getFormatVersionForFile(String file) { + Integer version = segmentFileVersionMap.get(file); + if (version != null) { + return version; + } + if (file.equals(segmentInfos.getSegmentsFileName())) { + return segmentInfos.getCommitLuceneVersion().major; + } + String segmentInfoFileName = RemoteStoreUtils.getSegmentName(file) + ".si"; + Integer siVersion = segmentFileVersionMap.get(segmentInfoFileName); + if (siVersion != null) { + return siVersion; + } + return org.apache.lucene.util.Version.LATEST.major; + } + + /** + * Serializes the actual SegmentInfos to bytes for the remote metadata file. + */ + @Override + public byte[] serialize() throws IOException { + ByteBuffersDataOutput byteBuffersIndexOutput = new ByteBuffersDataOutput(); + segmentInfos.write(new ByteBuffersIndexOutput(byteBuffersIndexOutput, "Snapshot of SegmentInfos", "SegmentInfos")); + return byteBuffersIndexOutput.toArrayCopy(); + } + + @Override + public Collection getFiles(boolean includeSegmentsFile) throws IOException { + return segmentInfos.files(includeSegmentsFile); + } + + private Map buildSegmentToLuceneVersionMap() { + Map segmentToLuceneVersion = new HashMap<>(); + for (SegmentCommitInfo segmentCommitInfo : segmentInfos) { + SegmentInfo info = segmentCommitInfo.info; + Set segFiles = info.files(); + for (String segFile : segFiles) { + segmentToLuceneVersion.put(segFile, info.getVersion().major); + } + } + return segmentToLuceneVersion; } } 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 2497d1ec7246b..36bbf3e89623b 100644 --- a/server/src/main/java/org/opensearch/index/shard/IndexShard.java +++ b/server/src/main/java/org/opensearch/index/shard/IndexShard.java @@ -147,6 +147,8 @@ import org.opensearch.index.engine.SegmentsStats; import org.opensearch.index.engine.dataformat.DataFormatRegistry; import org.opensearch.index.engine.exec.Indexer; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; +import org.opensearch.index.engine.exec.coord.SegmentInfosCatalogSnapshot; import org.opensearch.index.fielddata.FieldDataStats; import org.opensearch.index.fielddata.ShardFieldData; import org.opensearch.index.flush.FlushStats; @@ -181,6 +183,7 @@ import org.opensearch.index.shard.PrimaryReplicaSyncer.ResyncTask; import org.opensearch.index.similarity.SimilarityService; import org.opensearch.index.store.RemoteSegmentStoreDirectory; +import org.opensearch.index.store.RemoteSegmentStoreDirectory.UploadedSegmentMetadata; import org.opensearch.index.store.RemoteStoreFileDownloader; import org.opensearch.index.store.Store; import org.opensearch.index.store.Store.MetadataSnapshot; @@ -2030,6 +2033,39 @@ ReplicationCheckpoint computeReplicationCheckpoint(SegmentInfos segmentInfos) th return checkpoint; } + /** + * Compute the latest {@link ReplicationCheckpoint} from a CatalogSnapshot. + * This function fetches a metadata snapshot from the store that comes with an IO cost. + * We will reuse the existing stored checkpoint if it is at the same SI version. + * + * @param catalogSnapshot {@link CatalogSnapshot} infos to use to compute. + * @return {@link ReplicationCheckpoint} Checkpoint computed from the infos. + * @throws IOException When there is an error computing segment metadata from the store. + */ + ReplicationCheckpoint computeReplicationCheckpoint(CatalogSnapshot catalogSnapshot) throws IOException { + if (catalogSnapshot == null) { + return ReplicationCheckpoint.empty(shardId); + } + final ReplicationCheckpoint latestReplicationCheckpoint = getLatestReplicationCheckpoint(); + if (latestReplicationCheckpoint.getSegmentInfosVersion() == catalogSnapshot.getVersion() + && latestReplicationCheckpoint.getSegmentsGen() == catalogSnapshot.getGeneration() + && latestReplicationCheckpoint.getPrimaryTerm() == getOperationPrimaryTerm()) { + return latestReplicationCheckpoint; + } + final Map metadataMap = store.getSegmentMetadataMap(catalogSnapshot); + final ReplicationCheckpoint checkpoint = new ReplicationCheckpoint( + this.shardId, + getOperationPrimaryTerm(), + catalogSnapshot.getGeneration(), + catalogSnapshot.getVersion(), + metadataMap.values().stream().mapToLong(StoreFileMetadata::length).sum(), + getIndexer().config().getCodec().getName(), + metadataMap + ); + logger.trace("Recomputed ReplicationCheckpoint from CatalogSnapshot for shard {}", checkpoint); + return checkpoint; + } + public void publishReferencedSegments() throws IOException { assert referencedSegmentsPublisher != null; referencedSegmentsPublisher.publish(this, computeReferencedSegmentsCheckpoint()); @@ -5614,8 +5650,7 @@ public void syncSegmentsFromRemoteSegmentStore(boolean overrideLocal, final Runn // are uploaded to the remote segment store. RemoteSegmentMetadata remoteSegmentMetadata = remoteDirectory.init(); - Map uploadedSegments = remoteDirectory - .getSegmentsUploadedToRemoteStore() + Map uploadedSegments = remoteDirectory.getSegmentsUploadedToRemoteStore() .entrySet() .stream() .filter(entry -> entry.getKey().startsWith(IndexFileNames.SEGMENTS) == false) @@ -5691,8 +5726,7 @@ public void syncSegmentsFromGivenRemoteSegmentStore( remoteDirectory.init(); remoteStore.incRef(); } - Map uploadedSegments = sourceRemoteDirectory - .getSegmentsUploadedToRemoteStore(); + Map uploadedSegments = sourceRemoteDirectory.getSegmentsUploadedToRemoteStore(); store.incRef(); try { final Directory storeDirectory; @@ -5765,7 +5799,7 @@ private String copySegmentFiles( Directory storeDirectory, RemoteSegmentStoreDirectory sourceRemoteDirectory, RemoteSegmentStoreDirectory targetRemoteDirectory, - Map uploadedSegments, + Map uploadedSegments, boolean overrideLocal, final Runnable onFileSync ) throws IOException { @@ -5894,6 +5928,17 @@ public GatedCloseable getSegmentInfosSnapshot() { throw new IllegalStateException("Cannot request SegmentInfos directly on IndexShard"); } + /** + * Returns a reference-counted {@link CatalogSnapshot} for the current shard state. + * If a {@link DataFormatAwareEngine} is present, + * acquires from there. Otherwise wraps the SegmentInfos into a {@link SegmentInfosCatalogSnapshot}. + * + * @return a {@link GatedCloseable} wrapping the catalog snapshot + */ + public GatedCloseable getCatalogSnapshot() { + return getIndexer().acquireSnapshot(); + } + private TimeValue getRemoteTranslogUploadBufferInterval(Supplier clusterRemoteTranslogBufferIntervalSupplier) { assert Objects.nonNull(clusterRemoteTranslogBufferIntervalSupplier) : "remote translog buffer interval supplier is null"; if (indexSettings().isRemoteTranslogBufferIntervalExplicit()) { diff --git a/server/src/main/java/org/opensearch/index/shard/RemoteStoreRefreshListener.java b/server/src/main/java/org/opensearch/index/shard/RemoteStoreRefreshListener.java index 6b501c81b79d1..ecbcf83f30f45 100644 --- a/server/src/main/java/org/opensearch/index/shard/RemoteStoreRefreshListener.java +++ b/server/src/main/java/org/opensearch/index/shard/RemoteStoreRefreshListener.java @@ -28,8 +28,10 @@ import org.opensearch.index.engine.EngineBackedIndexer; import org.opensearch.index.engine.EngineException; import org.opensearch.index.engine.InternalEngine; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import org.opensearch.index.remote.RemoteSegmentTransferTracker; import org.opensearch.index.seqno.SequenceNumbers; +import org.opensearch.index.store.DataFormatAwareStoreDirectory; import org.opensearch.index.store.RemoteSegmentStoreDirectory; import org.opensearch.index.store.remote.metadata.RemoteSegmentMetadata; import org.opensearch.index.translog.Translog; @@ -137,8 +139,8 @@ protected void runAfterRefreshExactlyOnce(boolean didRefresh) { if (shouldSync(didRefresh, true) && isReadyForUpload()) { try { segmentTracker.updateLocalRefreshTimeAndSeqNo(); - try (GatedCloseable segmentInfosGatedCloseable = indexShard.getSegmentInfosSnapshot()) { - Collection localSegmentsPostRefresh = segmentInfosGatedCloseable.get().files(true); + try (GatedCloseable catalogSnapshotRef = indexShard.getCatalogSnapshot()) { + Collection localSegmentsPostRefresh = catalogSnapshotRef.get().getFiles(true); updateLocalSizeMapAndTracker(localSegmentsPostRefresh); } } catch (Throwable t) { @@ -206,8 +208,8 @@ private boolean shouldSync(boolean didRefresh, boolean skipPrimaryTermCheck) { * @return true iff all the local files are uploaded to remote store. */ boolean isRemoteSegmentStoreInSync() { - try (GatedCloseable segmentInfosGatedCloseable = indexShard.getSegmentInfosSnapshot()) { - return segmentInfosGatedCloseable.get().files(true).stream().allMatch(this::skipUpload); + try (GatedCloseable catalogSnapshotRef = indexShard.getCatalogSnapshot()) { + return catalogSnapshotRef.get().getFiles(true).stream().allMatch(this::skipUpload); } catch (Throwable throwable) { logger.error("Throwable thrown during isRemoteSegmentStoreInSync", throwable); } @@ -250,9 +252,9 @@ private boolean syncSegments() { remoteDirectory.deleteStaleSegmentsAsync(indexShard.getRemoteStoreSettings().getMinRemoteSegmentMetadataFiles()); } - try (GatedCloseable segmentInfosGatedCloseable = indexShard.getSegmentInfosSnapshot()) { - SegmentInfos segmentInfos = segmentInfosGatedCloseable.get(); - final ReplicationCheckpoint checkpoint = indexShard.computeReplicationCheckpoint(segmentInfos); + try (GatedCloseable catalogSnapshotRef = indexShard.getCatalogSnapshot()) { + CatalogSnapshot catalogSnapshot = catalogSnapshotRef.get(); + final ReplicationCheckpoint checkpoint = indexShard.computeReplicationCheckpoint(catalogSnapshot); if (checkpoint.getPrimaryTerm() != indexShard.getOperationPrimaryTerm()) { throw new IllegalStateException( String.format( @@ -266,7 +268,7 @@ private boolean syncSegments() { // Capture replication checkpoint before uploading the segments as upload can take some time and checkpoint can // move. long lastRefreshedCheckpoint = indexShard.getIndexer().lastRefreshedCheckpoint(); - Collection localSegmentsPostRefresh = segmentInfos.files(true); + Collection localSegmentsPostRefresh = catalogSnapshot.getFiles(true); // Create a map of file name to size and update the refresh segment tracker Map localSegmentsSizeMap = updateLocalSizeMapAndTracker(localSegmentsPostRefresh).entrySet() @@ -278,8 +280,8 @@ private boolean syncSegments() { public void onResponse(Void unused) { try { logger.debug("New segments upload successful"); - // Start metadata file upload in plaintext - uploadMetadata(localSegmentsPostRefresh, segmentInfos, checkpoint); + // Start metadata file upload + uploadMetadata(localSegmentsPostRefresh, catalogSnapshot, checkpoint); logger.debug("Metadata upload successful"); clearStaleFilesFromLocalSegmentChecksumMap(localSegmentsPostRefresh); onSuccessfulSegmentsSync( @@ -456,14 +458,17 @@ private boolean uploadedSegmentsMapExceedsThreshold() { return threshold != -1 && remoteDirectory.getSegmentsUploadedToRemoteStoreSize() > threshold; } - void uploadMetadata(Collection localSegmentsPostRefresh, SegmentInfos segmentInfos, ReplicationCheckpoint replicationCheckpoint) - throws IOException { + void uploadMetadata( + Collection localSegmentsPostRefresh, + CatalogSnapshot catalogSnapshot, + ReplicationCheckpoint replicationCheckpoint + ) throws IOException { final long maxSeqNo = indexShard.getIndexer().currentOngoingRefreshCheckpoint(); - SegmentInfos segmentInfosSnapshot = segmentInfos.clone(); - Map userData = segmentInfosSnapshot.getUserData(); + CatalogSnapshot catalogSnapshotCloned = catalogSnapshot.cloneNoAcquire(); + Map userData = new HashMap<>(catalogSnapshotCloned.getUserData()); userData.put(LOCAL_CHECKPOINT_KEY, String.valueOf(maxSeqNo)); userData.put(SequenceNumbers.MAX_SEQ_NO, Long.toString(maxSeqNo)); - segmentInfosSnapshot.setUserData(userData, false); + catalogSnapshotCloned.setUserData(userData, false); Translog.TranslogGeneration translogGeneration = indexShard.getIndexer().translogManager().getTranslogGeneration(); if (translogGeneration == null) { @@ -472,7 +477,7 @@ void uploadMetadata(Collection localSegmentsPostRefresh, SegmentInfos se long translogFileGeneration = translogGeneration.translogFileGeneration; remoteDirectory.uploadMetadata( localSegmentsPostRefresh, - segmentInfosSnapshot, + catalogSnapshotCloned, storeDirectory, translogFileGeneration, replicationCheckpoint, @@ -506,6 +511,15 @@ private boolean skipUpload(String file) { private String getChecksumOfLocalFile(String file) throws IOException { if (!localSegmentChecksumMap.containsKey(file)) { + if (indexShard.indexSettings().isPluggableDataFormatEnabled()) { + DataFormatAwareStoreDirectory dfasd = DataFormatAwareStoreDirectory.unwrap(storeDirectory); + if (dfasd == null) { + throw new IllegalStateException("DataFormatAwareStoreDirectory expected when pluggable data format is enabled"); + } + String checksum = dfasd.calculateUploadChecksum(file); + localSegmentChecksumMap.put(file, checksum); + return checksum; + } try (IndexInput indexInput = storeDirectory.openInput(file, IOContext.READONCE)) { String checksum = Long.toString(CodecUtil.retrieveChecksum(indexInput)); localSegmentChecksumMap.put(file, checksum); diff --git a/server/src/main/java/org/opensearch/index/store/DataFormatAwareStoreDirectory.java b/server/src/main/java/org/opensearch/index/store/DataFormatAwareStoreDirectory.java new file mode 100644 index 0000000000000..24065799c537e --- /dev/null +++ b/server/src/main/java/org/opensearch/index/store/DataFormatAwareStoreDirectory.java @@ -0,0 +1,298 @@ +/* + * 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.store; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FilterDirectory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.IndexOutput; +import org.opensearch.common.annotation.PublicApi; +import org.opensearch.index.IndexSettings; +import org.opensearch.index.engine.dataformat.DataFormatDescriptor; +import org.opensearch.index.engine.dataformat.DataFormatRegistry; +import org.opensearch.index.shard.ShardPath; +import org.opensearch.index.store.checksum.GenericCRC32ChecksumHandler; +import org.opensearch.index.store.checksum.LuceneChecksumHandler; + +import java.io.IOException; +import java.util.Collection; +import java.util.HashMap; +import java.util.Locale; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; + +/** + * Format-aware directory that extends {@link FilterDirectory} and wraps a {@link SubdirectoryAwareDirectory}. + * + *

This directory adds data format awareness on top of the subdirectory path routing + * provided by {@link SubdirectoryAwareDirectory}. It understands that files in different + * subdirectories may belong to different data formats (Lucene, Parquet, Arrow, etc.) and + * provides format-specific operations, most notably checksum calculation.

+ * + *

Delegated to SubdirectoryAwareDirectory:

+ *
    + *
  • Path routing: plain filenames → index/, prefixed filenames → subdirectories
    • + *
  • File operations: openInput, createOutput, deleteFile, fileLength, listAll, rename, sync
    • + *
+ * + *

Added by DataFormatAwareStoreDirectory:

+ *
    + *
  • FileMetadata support: parse file identifier strings into FileMetadata objects
    • + *
  • Format-aware checksum: Lucene files → CodecUtil footer, others → CRC32 full-file
    • + *
  • Dual API: callers can use String or FileMetadata for all operations
    • + *
+ * + *

File naming convention:

+ * 
+ *   Lucene:   "_0.cfs"                → stored in <shard>/index/_0.cfs
+ *   Parquet:  "parquet/_0_1.parquet"  → stored in <shard>/parquet/_0_1.parquet
+ *   Arrow:    "arrow/_0_1.arrow"      → stored in <shard>/arrow/_0_1.arrow
+ *
+ * + *

Checksum strategy:

+ *
    + *
  • Lucene/index files: {@code CodecUtil.retrieveChecksum()} — reads checksum from codec footer (fast, O(1))
    • + *
  • Non-Lucene files: Full-file CRC32 scan — computes CRC32 over all bytes (generic, O(n))
    • + *
+ * + * @opensearch.api + */ +@PublicApi(since = "3.0.0") +public class DataFormatAwareStoreDirectory extends FilterDirectory { + + private static final Logger logger = LogManager.getLogger(DataFormatAwareStoreDirectory.class); + + private static final String DEFAULT_FORMAT = "lucene"; + + private static final Set INDEX_DIRECTORY_FORMATS = Set.of("lucene", "metadata"); + + private final ShardPath shardPath; + private final Map checksumStrategies; + private static final FormatChecksumStrategy DEFAULT_CHECKSUM_STRATEGY = new GenericCRC32ChecksumHandler(); + + /** + * Constructs a DataFormatAwareStoreDirectory with a {@link DataFormatRegistry} for format-aware + * checksum calculation and other format-specific operations. + * + * @param delegate the underlying FSDirectory (typically for <shard>/index/) + * @param shardPath the shard path for resolving subdirectories + * @param dataFormatRegistry registry providing format-specific checksum handlers + */ + public DataFormatAwareStoreDirectory( + IndexSettings indexSettings, + Directory delegate, + ShardPath shardPath, + DataFormatRegistry dataFormatRegistry + ) { + super(new SubdirectoryAwareDirectory(delegate, shardPath)); + this.shardPath = shardPath; + Map descriptors = dataFormatRegistry.getFormatDescriptors(indexSettings); + this.checksumStrategies = new HashMap<>(); + for (Map.Entry entry : descriptors.entrySet()) { + this.checksumStrategies.put(entry.getKey(), entry.getValue().getChecksumStrategy()); + } + this.checksumStrategies.put(DEFAULT_FORMAT, new LuceneChecksumHandler()); + + logger.debug( + "Created DataFormatAwareStoreDirectory for shard {} with checksum strategies for formats: {}", + shardPath.getShardId(), + checksumStrategies.keySet() + ); + } + + /** + * Walks the {@link FilterDirectory} wrapping chain to find a {@link DataFormatAwareStoreDirectory}. + * This is needed because the directory may be wrapped in {@link ByteSizeCachingDirectory} and + * {@link Store.StoreDirectory}, so a direct {@code instanceof} check on the outermost directory + * would fail. + * + * @param dir the directory to unwrap (may be null) + * @return the DataFormatAwareStoreDirectory found in the chain, or null if not present + */ + public static DataFormatAwareStoreDirectory unwrap(Directory dir) { + while (dir != null) { + if (dir instanceof DataFormatAwareStoreDirectory) { + return (DataFormatAwareStoreDirectory) dir; + } + if (dir instanceof FilterDirectory) { + dir = ((FilterDirectory) dir).getDelegate(); + } else { + return null; + } + } + return null; + } + + private String resolveFileName(String fileName) { + if (fileName.contains(FileMetadata.DELIMITER)) { + FileMetadata fm = new FileMetadata(fileName); + fileName = toFileIdentifier(fm); + } + return fileName; + } + + @Override + public IndexInput openInput(String name, IOContext context) throws IOException { + return in.openInput(resolveFileName(name), context); + } + + @Override + public IndexOutput createOutput(String name, IOContext context) throws IOException { + return in.createOutput(resolveFileName(name), context); + } + + @Override + public void deleteFile(String name) throws IOException { + in.deleteFile(resolveFileName(name)); + } + + @Override + public long fileLength(String name) throws IOException { + return in.fileLength(resolveFileName(name)); + } + + @Override + public void sync(Collection names) throws IOException { + in.sync(names.stream().map(this::resolveFileName).collect(Collectors.toList())); + } + + @Override + public void rename(String source, String dest) throws IOException { + in.rename(resolveFileName(source), resolveFileName(dest)); + } + + @Override + public String[] listAll() throws IOException { + String[] allFiles = in.listAll(); + for (int i = 0; i < allFiles.length; i++) { + // Normalize OS-dependent separators (e.g., "\" on Windows) to "/" before parsing, + // since SubdirectoryAwareDirectory.listAll() returns Path.toString() which uses + // the OS separator, but FileMetadata expects "/" as the format/file delimiter. + String normalized = allFiles[i].replace( + org.opensearch.common.io.PathUtils.getDefaultFileSystem().getSeparator().charAt(0), + '/' + ); + FileMetadata fm = toFileMetadata(normalized); + allFiles[i] = isDefaultFormat(fm.dataFormat()) ? fm.file() : fm.serialize(); + } + return allFiles; + } + + // ═══════════════════════════════════════════════════════════════ + // FileMetadata parsing and conversion + // ═══════════════════════════════════════════════════════════════ + + /** + * Parses a file identifier string into a {@link FileMetadata} object. + * Uses the same "format/file" convention as {@link FileMetadata} (e.g., "parquet/_0.pqt"). + * Plain filenames without a "/" prefix default to the lucene format. + * + * @param fileIdentifier the file path string (with optional format prefix separated by '/') + * @return FileMetadata with parsed dataFormat and filename + */ + public static FileMetadata toFileMetadata(String fileIdentifier) { + return new FileMetadata(fileIdentifier); + } + + /** + * Converts a {@link FileMetadata} object back to a file identifier string. + * + * @param fm the FileMetadata to convert + * @return file identifier string suitable for Directory operations + */ + public static String toFileIdentifier(FileMetadata fm) { + String format = fm.dataFormat(); + if (isDefaultFormat(format)) { + return fm.file(); + } + return format + "/" + fm.file(); + } + + // ═══════════════════════════════════════════════════════════════ + // Format-Aware Checksum Calculation + // ═══════════════════════════════════════════════════════════════ + + public long calculateChecksum(String name) throws IOException { + FileMetadata fm = toFileMetadata(name); + return calculateChecksum(fm); + } + + /** + * Calculates checksum using the format-specific {@link FormatChecksumStrategy}. + * Supports pre-computed checksums (O(1) for formats that register them during write) + * and falls back to file-based computation for formats that don't. + */ + private long calculateChecksum(FileMetadata fm) throws IOException { + String fileIdentifier = toFileIdentifier(fm); + FormatChecksumStrategy strategy = checksumStrategies.getOrDefault(fm.dataFormat(), DEFAULT_CHECKSUM_STRATEGY); + return strategy.computeChecksum(this, fileIdentifier); + } + + /** + * Calculates a checksum suitable for upload verification. + * Public API used by RemoteSegmentStoreDirectory. + */ + public String calculateUploadChecksum(String name) throws IOException { + return Long.toString(calculateChecksum(name)); + } + + /** + * Registers a {@link FormatChecksumStrategy} for a data format. + * Overrides any existing strategy + * + *

Use this to register strategies that support pre-computed checksums (e.g., + * {@link PrecomputedChecksumStrategy} for Parquet files whose CRC32 is computed + * during write by the Rust writer). + * + * @param format the data format name (e.g., "parquet") + * @param strategy the checksum strategy to use for this format + */ + public void registerChecksumStrategy(String format, FormatChecksumStrategy strategy) { + if (format != null && strategy != null) { + checksumStrategies.put(format, strategy); + logger.debug("Registered FormatChecksumStrategy for format [{}]", format); + } + } + + /** + * Returns the checksum strategy for the given format, or {@code null} if none is registered. + * Engines use this to share the directory's strategy instance so that pre-computed + * checksums registered during write are visible to the upload path. + * + * @param format the data format name (e.g., "parquet") + * @return the strategy, or null if not found + */ + public FormatChecksumStrategy getChecksumStrategy(String format) { + return checksumStrategies.get(format); + } + + public IndexOutput createOutput(FileMetadata fm, IOContext context) throws IOException { + return createOutput(toFileIdentifier(fm), context); + } + + public String getDataFormat(String fileIdentifier) { + return toFileMetadata(fileIdentifier).dataFormat(); + } + + public ShardPath getShardPath() { + return shardPath; + } + + // ═══════════════════════════════════════════════════════════════ + // Private Helpers + // ═══════════════════════════════════════════════════════════════ + + private static boolean isDefaultFormat(String format) { + return format == null || format.isEmpty() || INDEX_DIRECTORY_FORMATS.contains(format.toLowerCase(Locale.ROOT)); + } +} diff --git a/server/src/main/java/org/opensearch/index/store/DataFormatAwareStoreDirectoryFactory.java b/server/src/main/java/org/opensearch/index/store/DataFormatAwareStoreDirectoryFactory.java new file mode 100644 index 0000000000000..b633a00ca67eb --- /dev/null +++ b/server/src/main/java/org/opensearch/index/store/DataFormatAwareStoreDirectoryFactory.java @@ -0,0 +1,60 @@ +/* + * 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.store; + +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.core.index.shard.ShardId; +import org.opensearch.index.IndexSettings; +import org.opensearch.index.engine.dataformat.DataFormatRegistry; +import org.opensearch.index.shard.ShardPath; +import org.opensearch.plugins.IndexStorePlugin; + +import java.io.IOException; + +/** + * Factory interface for creating DataFormatAwareStoreDirectory instances. + * This interface follows the existing IndexStorePlugin pattern to provide + * a centralized way to create composite directories with format discovery. + * + *

Following the same delegation pattern as {@link IndexStorePlugin.CompositeDirectoryFactory}, + * this factory accepts a {@link IndexStorePlugin.DirectoryFactory} to delegate local directory + * creation rather than hardcoding a specific directory implementation. + * + * @opensearch.experimental + */ +@ExperimentalApi +@FunctionalInterface +public interface DataFormatAwareStoreDirectoryFactory { + + /** + * Creates a new DataFormatAwareStoreDirectory per shard with automatic format discovery. + *

+ * The factory will: + * - Delegate local directory creation to the provided localDirectoryFactory + * - Use DataFormatRegistry to discover available data format plugins + * - Create format-specific directories for each discovered format + * - Provide fallback behavior if no plugins are found + * - Handle errors gracefully with proper logging + * + * @param indexSettings the shard's index settings containing configuration + * @param shardId the shard identifier + * @param shardPath the path the shard is using for file storage + * @param localDirectoryFactory the factory for creating the underlying local directory, respecting index store type configuration + * @param dataFormatRegistry registry of available data format plugins + * @return a new DataFormatAwareStoreDirectory instance supporting all discovered formats + * @throws IOException if directory creation fails or resources cannot be allocated + */ + DataFormatAwareStoreDirectory newDataFormatAwareStoreDirectory( + IndexSettings indexSettings, + ShardId shardId, + ShardPath shardPath, + IndexStorePlugin.DirectoryFactory localDirectoryFactory, + DataFormatRegistry dataFormatRegistry + ) throws IOException; +} diff --git a/server/src/main/java/org/opensearch/index/store/DefaultDataFormatAwareStoreDirectoryFactory.java b/server/src/main/java/org/opensearch/index/store/DefaultDataFormatAwareStoreDirectoryFactory.java new file mode 100644 index 0000000000000..8e32942f5676d --- /dev/null +++ b/server/src/main/java/org/opensearch/index/store/DefaultDataFormatAwareStoreDirectoryFactory.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.index.store; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.apache.lucene.store.Directory; +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.core.index.shard.ShardId; +import org.opensearch.index.IndexSettings; +import org.opensearch.index.engine.dataformat.DataFormatRegistry; +import org.opensearch.index.shard.ShardPath; +import org.opensearch.plugins.IndexStorePlugin; + +import java.io.IOException; +import java.util.Locale; + +/** + * Default implementation of DataFormatAwareStoreDirectoryFactory that provides + * plugin-based format discovery and fallback behavior. + * + *

Delegates local directory creation to the provided {@link IndexStorePlugin.DirectoryFactory}, + * following the same pattern as {@link org.opensearch.index.store.DefaultCompositeDirectoryFactory}. + * + * @opensearch.experimental + */ +@ExperimentalApi() +public class DefaultDataFormatAwareStoreDirectoryFactory implements DataFormatAwareStoreDirectoryFactory { + + private static final Logger logger = LogManager.getLogger(DefaultDataFormatAwareStoreDirectoryFactory.class); + + /** + * Creates a new DataFormatAwareStoreDirectory with plugin-based format discovery. + * + * @param indexSettings the shard's index settings + * @param shardId the shard identifier + * @param shardPath the path the shard is using + * @param localDirectoryFactory the factory for creating the underlying local directory + * @param dataFormatRegistry registry of available data format plugins + * @return a new DataFormatAwareStoreDirectory instance + * @throws IOException if directory creation fails + */ + @Override + public DataFormatAwareStoreDirectory newDataFormatAwareStoreDirectory( + IndexSettings indexSettings, + ShardId shardId, + ShardPath shardPath, + IndexStorePlugin.DirectoryFactory localDirectoryFactory, + DataFormatRegistry dataFormatRegistry + ) throws IOException { + + if (logger.isDebugEnabled()) { + logger.debug( + "Creating DataFormatAwareStoreDirectory for shard: {} at path: {}", + shardPath.getShardId(), + shardPath.getDataPath() + ); + } + + try { + // Delegate local directory creation to the configured DirectoryFactory + Directory delegate = localDirectoryFactory.newDirectory(indexSettings, shardPath); + + DataFormatAwareStoreDirectory directory = new DataFormatAwareStoreDirectory( + indexSettings, + delegate, + shardPath, + dataFormatRegistry + ); + + if (logger.isDebugEnabled()) { + logger.debug( + "Successfully created DataFormatAwareStoreDirectory for shard: {} with registered formats: {}", + shardPath.getShardId(), + dataFormatRegistry.getRegisteredFormats() + ); + } + + return directory; + + } catch (Exception e) { + logger.error( + () -> new org.apache.logging.log4j.message.ParameterizedMessage( + "Failed to create DataFormatAwareStoreDirectory for shard: {}", + shardPath.getShardId() + ), + e + ); + throw new IOException( + String.format( + Locale.ROOT, + "Failed to create DataFormatAwareStoreDirectory for shard %s: %s", + shardPath.getShardId(), + e.getMessage() + ), + e + ); + } + } +} diff --git a/server/src/main/java/org/opensearch/index/store/FileMetadata.java b/server/src/main/java/org/opensearch/index/store/FileMetadata.java new file mode 100644 index 0000000000000..980e27c5d9014 --- /dev/null +++ b/server/src/main/java/org/opensearch/index/store/FileMetadata.java @@ -0,0 +1,152 @@ +/* + * 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.store; + +import org.opensearch.common.annotation.ExperimentalApi; + +import java.util.Objects; + +/** + * Represents metadata for a file in the index, including its data format and filename. + * Files can be in different formats (e.g., "lucene", "metadata") and this class provides + * a unified way to represent and serialize file information across the system. + * + * @opensearch.experimental + */ +@ExperimentalApi +public class FileMetadata { + + /** + * Delimiter used to separate data format and filename in serialized form. + * Uses "/" to match the subdirectory convention (e.g., "parquet/_0.pqt"). + */ + public static final String DELIMITER = "/"; + private static final String DEFAULT_FORMAT = "lucene"; + private static final String METADATA_KEY = "metadata"; + + private final String file; + private final String dataFormat; + + /** + * Constructs a FileMetadata with explicit data format and filename. + * + * @param dataFormat the data format identifier (e.g., "lucene", "metadata") + * @param file the filename + */ + public FileMetadata(String dataFormat, String file) { + this.file = file; + this.dataFormat = dataFormat; + } + + /** + * Constructs a FileMetadata by parsing a serialized data-format-aware filename. + * The format is "format/file" (e.g., "parquet/_0.pqt"). If no delimiter is present, + * files starting with "metadata" are treated as metadata format, otherwise defaults to "lucene". + * + * @param dataFormatAwareFile the serialized filename with optional data format prefix + */ + public FileMetadata(String dataFormatAwareFile) { + int slash = dataFormatAwareFile.indexOf(DELIMITER); + if (slash >= 0) { + this.dataFormat = dataFormatAwareFile.substring(0, slash); + this.file = dataFormatAwareFile.substring(slash + 1); + } else if (dataFormatAwareFile.startsWith(METADATA_KEY)) { + this.dataFormat = METADATA_KEY; + this.file = dataFormatAwareFile; + } else { + this.dataFormat = DEFAULT_FORMAT; + this.file = dataFormatAwareFile; + } + } + + /** + * Serializes a data format and filename into a format-aware string without creating an intermediate object. + * For the default lucene format, returns just the filename (no prefix). + * + * @param dataFormat the data format identifier (e.g., "lucene", "parquet") + * @param file the filename + * @return the serialized representation (e.g., "parquet/_0.parquet" or "_0.si" for lucene) + */ + public static String serialize(String dataFormat, String file) { + if (DEFAULT_FORMAT.equals(dataFormat)) { + return file; + } + return dataFormat + DELIMITER + file; + } + + /** + * Extracts the plain filename from a serialized format-aware filename without creating an intermediate object. + * + * @param serialized the serialized filename (e.g., "parquet/_0.parquet" or "_0.si") + * @return the plain filename (e.g., "_0.parquet" or "_0.si") + */ + public static String parseFile(String serialized) { + int slash = serialized.indexOf(DELIMITER); + return slash >= 0 ? serialized.substring(slash + 1) : serialized; + } + + /** + * Extracts the data format from a serialized format-aware filename without creating an intermediate object. + * + * @param serialized the serialized filename (e.g., "parquet/_0.parquet" or "_0.si") + * @return the data format (e.g., "parquet", "lucene", or "metadata") + */ + public static String parseDataFormat(String serialized) { + int slash = serialized.indexOf(DELIMITER); + if (slash >= 0) { + return serialized.substring(0, slash); + } + return serialized.startsWith(METADATA_KEY) ? METADATA_KEY : DEFAULT_FORMAT; + } + + /** + * Serializes this FileMetadata to a string in the format "format/file". + * For the default lucene format, returns just the filename (no prefix). + * + * @return the serialized representation + */ + public String serialize() { + return serialize(dataFormat, file); + } + + @Override + public String toString() { + return serialize(); + } + + /** + * Returns the filename. + * + * @return the filename + */ + public String file() { + return file; + } + + /** + * Returns the data format identifier. + * + * @return the data format (e.g., "lucene", "metadata") + */ + public String dataFormat() { + return dataFormat; + } + + @Override + public boolean equals(Object o) { + if (o == null || getClass() != o.getClass()) return false; + FileMetadata that = (FileMetadata) o; + return Objects.equals(file, that.file) && Objects.equals(dataFormat, that.dataFormat); + } + + @Override + public int hashCode() { + return Objects.hash(file, dataFormat); + } +} diff --git a/server/src/main/java/org/opensearch/index/store/FormatChecksumStrategy.java b/server/src/main/java/org/opensearch/index/store/FormatChecksumStrategy.java new file mode 100644 index 0000000000000..b48ec89a9688e --- /dev/null +++ b/server/src/main/java/org/opensearch/index/store/FormatChecksumStrategy.java @@ -0,0 +1,71 @@ +/* + * 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.store; + +import org.apache.lucene.store.Directory; +import org.opensearch.common.annotation.ExperimentalApi; + +import java.io.IOException; + +/** + * Strategy for computing file checksums, with support for pre-computed values. + * + *

Each data format provides its own strategy: + *

    + *
  • Lucene: reads the codec footer (O(1))
    • + *
  • Parquet: uses pre-computed CRC32 from the Rust writer (O(1)), + * falls back to full-file scan if not available
    • + *
  • Default: full-file CRC32 scan (O(n))
    • + *
+ * + *

Pre-computed checksums are registered via {@link #registerChecksum(String, long, long)} + * during the flush/write path, then consumed by the upload path via + * {@link #computeChecksum(Directory, String)}. + * + *

The {@code writerGeneration} parameter in {@link #registerChecksum} ensures that + * if a filename is reused across generations (e.g., after merge or rewrite), the cache + * entry is always from the latest write. The generation acts as a version stamp — a + * stale entry from an older generation is overwritten, never served. + * + * @opensearch.experimental + */ +@ExperimentalApi +public interface FormatChecksumStrategy { + + /** + * Computes or retrieves the checksum for the given file. + * + * @param dir the directory containing the file + * @param fileName the file name (local name, not format-prefixed) + * @return the checksum value + * @throws IOException if checksum computation fails + */ + long computeChecksum(Directory dir, String fileName) throws IOException; + + /** + * Registers a pre-computed checksum for a file at a specific writer generation. + * Called during the write/flush path when the checksum is known. + * + *

The generation parameter ensures cache correctness: if the same filename + * is written by a later generation, the new checksum replaces the old one. + * Implementations should store the generation alongside the checksum and + * reject lookups for stale generations. + * + * @param fileName the file name + * @param checksum the pre-computed checksum + * @param writerGeneration the writer generation that produced this file + */ + default void registerChecksum(String fileName, long checksum, long writerGeneration) {} + + /** + * Clears all cached checksums. Called during cleanup/close. + */ + default void clearChecksums() {} + +} diff --git a/server/src/main/java/org/opensearch/index/store/PrecomputedChecksumStrategy.java b/server/src/main/java/org/opensearch/index/store/PrecomputedChecksumStrategy.java new file mode 100644 index 0000000000000..293e31d160e0f --- /dev/null +++ b/server/src/main/java/org/opensearch/index/store/PrecomputedChecksumStrategy.java @@ -0,0 +1,100 @@ +/* + * 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.store; + +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.opensearch.common.annotation.ExperimentalApi; + +import java.io.IOException; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.zip.CRC32; + +/** + * Checksum strategy that uses pre-computed checksums when available, + * falling back to full-file CRC32 scan. + * + *

The write path registers checksums via {@link #registerChecksum(String, long, long)}. + * The upload path retrieves them via {@link #computeChecksum(Directory, String)} — O(1) + * if pre-computed, O(n) fallback otherwise. + * + *

Each cache entry stores the checksum alongside the writer generation that produced it. + * This ensures that if a filename is reused (e.g., after merge), the cache always serves + * the checksum from the latest write, not a stale value from an earlier generation. + * + * @opensearch.experimental + */ +@ExperimentalApi +public class PrecomputedChecksumStrategy implements FormatChecksumStrategy { + + /** Cache entry: checksum + the generation that produced it. */ + private record CacheEntry(long checksum, long generation) { + } + + private final Map checksumCache = new ConcurrentHashMap<>(); + + @Override + public long computeChecksum(Directory dir, String fileName) throws IOException { + CacheEntry entry = checksumCache.get(fileName); + if (entry != null) { + return entry.checksum(); + } + // Fallback: full-file CRC32 scan + return computeFullFileCrc32(dir, fileName); + } + + @Override + public void registerChecksum(String fileName, long checksum, long writerGeneration) { + if (fileName != null && checksum != 0) { + checksumCache.compute(fileName, (key, existing) -> { + // Only overwrite if the new generation is >= the existing one. + // This prevents a race where an older generation's late registration + // overwrites a newer generation's checksum. + if (existing == null || writerGeneration >= existing.generation()) { + return new CacheEntry(checksum, writerGeneration); + } + return existing; + }); + } + } + + @Override + public void clearChecksums() { + checksumCache.clear(); + } + + /** + * Removes a single checksum entry after it has been consumed (e.g., after successful upload). + * Prevents unbounded cache growth over the shard's lifetime. + * + * @param fileName the file whose checksum should be evicted + */ + public void evictChecksum(String fileName) { + if (fileName != null) { + checksumCache.remove(fileName); + } + } + + private static long computeFullFileCrc32(Directory dir, String fileName) throws IOException { + CRC32 crc32 = new CRC32(); + byte[] buffer = new byte[64 * 1024]; + try (IndexInput input = dir.openInput(fileName, IOContext.READONCE)) { + long remaining = input.length(); + while (remaining > 0) { + int toRead = (int) Math.min(buffer.length, remaining); + input.readBytes(buffer, 0, toRead); + crc32.update(buffer, 0, toRead); + remaining -= toRead; + } + } + return crc32.getValue(); + } +} diff --git a/server/src/main/java/org/opensearch/index/store/RemoteDirectory.java b/server/src/main/java/org/opensearch/index/store/RemoteDirectory.java index e434d5813311b..ff434ab33ad3b 100644 --- a/server/src/main/java/org/opensearch/index/store/RemoteDirectory.java +++ b/server/src/main/java/org/opensearch/index/store/RemoteDirectory.java @@ -33,6 +33,7 @@ import org.opensearch.core.action.ActionListener; import org.opensearch.core.common.unit.ByteSizeUnit; import org.opensearch.index.store.exception.ChecksumCombinationException; +import org.opensearch.index.store.remote.FormatBlobRouter; import java.io.FileNotFoundException; import java.io.IOException; @@ -43,6 +44,7 @@ import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Set; import java.util.concurrent.CountDownLatch; import java.util.concurrent.atomic.AtomicReference; @@ -75,7 +77,7 @@ public class RemoteDirectory extends Directory { * Map containing the mapping of segment files that are pending download as part of the pre-copy (warm) phase of * {@link org.opensearch.index.engine.MergedSegmentWarmer}. The key is the local filename and value is the remote filename. */ - final Map pendingDownloadMergedSegments; + protected final Map pendingDownloadMergedSegments; /** * Number of bytes in the segment file to store checksum @@ -382,6 +384,10 @@ public void delete() throws IOException { blobContainer.delete(); } + public Optional getFormatBlobRouter() { + return Optional.empty(); + } + public boolean copyFrom( Directory from, String src, @@ -403,7 +409,7 @@ public boolean copyFrom( return false; } - private void uploadBlob( + protected void uploadBlob( Directory from, String src, String remoteFileName, @@ -412,6 +418,20 @@ private void uploadBlob( ActionListener listener, boolean lowPriorityUpload, CryptoMetadata cryptoMetadata + ) throws Exception { + uploadBlob(from, src, remoteFileName, ioContext, postUploadRunner, listener, lowPriorityUpload, cryptoMetadata, blobContainer); + } + + protected void uploadBlob( + Directory from, + String src, + String remoteFileName, + IOContext ioContext, + Runnable postUploadRunner, + ActionListener listener, + boolean lowPriorityUpload, + CryptoMetadata cryptoMetadata, + BlobContainer targetBlobContainer ) throws Exception { assert ioContext != IOContext.READONCE : "Remote upload will fail with IoContext.READONCE"; long expectedChecksum = calculateChecksumOfChecksum(from, src); @@ -420,7 +440,7 @@ private void uploadBlob( try { contentLength = indexInput.length(); boolean remoteIntegrityEnabled = false; - if (getBlobContainer() instanceof AsyncMultiStreamBlobContainer asyncContainer) { + if (targetBlobContainer instanceof AsyncMultiStreamBlobContainer asyncContainer) { remoteIntegrityEnabled = asyncContainer.remoteIntegrityCheckSupported(); } lowPriorityUpload = lowPriorityUpload || contentLength > ByteSizeUnit.GB.toBytes(15); @@ -488,7 +508,7 @@ private void uploadBlob( }); WriteContext writeContext = remoteTransferContainer.createWriteContext(); - ((AsyncMultiStreamBlobContainer) blobContainer).asyncBlobUpload(writeContext, completionListener); + ((AsyncMultiStreamBlobContainer) targetBlobContainer).asyncBlobUpload(writeContext, completionListener); } catch (Exception e) { logger.warn("Exception while calling asyncBlobUpload, closing IndexInput to avoid leak"); indexInput.close(); @@ -496,7 +516,7 @@ private void uploadBlob( } } - private long calculateChecksumOfChecksum(Directory directory, String file) throws IOException { + protected long calculateChecksumOfChecksum(Directory directory, String file) throws IOException { try (IndexInput indexInput = directory.openInput(file, IOContext.READONCE)) { try { return checksumOfChecksum(indexInput, SEGMENT_CHECKSUM_BYTES); diff --git a/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java b/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java index ab678c0ffe2f4..801692b2b7da8 100644 --- a/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java +++ b/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java @@ -35,12 +35,14 @@ import org.opensearch.common.lucene.store.ByteArrayIndexInput; import org.opensearch.core.action.ActionListener; import org.opensearch.core.index.shard.ShardId; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import org.opensearch.index.remote.RemoteStorePathStrategy; import org.opensearch.index.remote.RemoteStoreUtils; import org.opensearch.index.store.lockmanager.FileLockInfo; import org.opensearch.index.store.lockmanager.RemoteStoreCommitLevelLockManager; import org.opensearch.index.store.lockmanager.RemoteStoreLockManager; import org.opensearch.index.store.lockmanager.RemoteStoreMetadataLockManager; +import org.opensearch.index.store.remote.FormatBlobRouter; import org.opensearch.index.store.remote.metadata.RemoteSegmentMetadata; import org.opensearch.index.store.remote.metadata.RemoteSegmentMetadataHandlerFactory; import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint; @@ -111,9 +113,21 @@ public final class RemoteSegmentStoreDirectory extends FilterDirectory implement * Keeps track of local segment filename to uploaded filename along with other attributes like checksum. * This map acts as a cache layer for uploaded segment filenames which helps avoid calling listAll() each time. * It is important to initialize this map on creation of RemoteSegmentStoreDirectory and update it on each upload and delete. + * + *

IMPORTANT: All mutations to this map must go through the encapsulated APIs + * ({@link #replaceUploadedSegments}, {@link #addUploadedSegment}, {@link #removeUploadedSegment}) + * to keep the format cache in FormatBlobRouter in sync. */ private Map segmentsUploadedToRemoteStore; + /** + * Format blob router for managing blob key → format reverse lookup cache. + * Non-null only when remoteDataDirectory is a DataFormatAwareRemoteDirectory. + * When null, format registration calls are no-ops (single-format index). + */ + @Nullable + private final FormatBlobRouter formatBlobRouter; + private static final VersionedCodecStreamWrapper metadataStreamWrapper = new VersionedCodecStreamWrapper<>( new RemoteSegmentMetadataHandlerFactory(), RemoteSegmentMetadata.VERSION_ONE, @@ -162,6 +176,7 @@ public RemoteSegmentStoreDirectory( this.metadataFilePinnedTimestampMap = new HashMap<>(); this.logger = Loggers.getLogger(getClass(), shardId); this.pendingDownloadMergedSegments = pendingDownloadMergedSegments; + this.formatBlobRouter = remoteDataDirectory.getFormatBlobRouter().orElse(null); init(); } @@ -177,9 +192,9 @@ public RemoteSegmentMetadata init() throws IOException { logger.debug("Start initialisation of remote segment metadata"); RemoteSegmentMetadata remoteSegmentMetadata = readLatestMetadataFile(); if (remoteSegmentMetadata != null) { - this.segmentsUploadedToRemoteStore = new ConcurrentHashMap<>(remoteSegmentMetadata.getMetadata()); + replaceUploadedSegments(remoteSegmentMetadata.getMetadata()); } else { - this.segmentsUploadedToRemoteStore = new ConcurrentHashMap<>(); + replaceUploadedSegments(Collections.emptyMap()); } logger.debug("Initialisation of remote segment metadata completed"); return remoteSegmentMetadata; @@ -198,9 +213,9 @@ public RemoteSegmentMetadata initializeToSpecificCommit(long primaryTerm, long c String metadataFile = ((RemoteStoreMetadataLockManager) mdLockManager).fetchLockedMetadataFile(metadataFilePrefix, acquirerId); RemoteSegmentMetadata remoteSegmentMetadata = readMetadataFile(metadataFile); if (remoteSegmentMetadata != null) { - this.segmentsUploadedToRemoteStore = new ConcurrentHashMap<>(remoteSegmentMetadata.getMetadata()); + replaceUploadedSegments(remoteSegmentMetadata.getMetadata()); } else { - this.segmentsUploadedToRemoteStore = new ConcurrentHashMap<>(); + replaceUploadedSegments(Collections.emptyMap()); } return remoteSegmentMetadata; } @@ -235,9 +250,9 @@ public RemoteSegmentMetadata initializeToSpecificTimestamp(long timestamp) throw String metadataFile = lockedMetadataFiles.iterator().next(); RemoteSegmentMetadata remoteSegmentMetadata = readMetadataFile(metadataFile); if (remoteSegmentMetadata != null) { - this.segmentsUploadedToRemoteStore = new ConcurrentHashMap<>(remoteSegmentMetadata.getMetadata()); + replaceUploadedSegments(remoteSegmentMetadata.getMetadata()); } else { - this.segmentsUploadedToRemoteStore = new ConcurrentHashMap<>(); + replaceUploadedSegments(Collections.emptyMap()); } return remoteSegmentMetadata; } @@ -510,8 +525,10 @@ public String[] listAll() throws IOException { public void deleteFile(String name) throws IOException { String remoteFilename = getExistingRemoteFilename(name); if (remoteFilename != null) { + // Step 1: delete from remote (format cache entry still available for routing) remoteDataDirectory.deleteFile(remoteFilename); - segmentsUploadedToRemoteStore.remove(name); + // Step 2: cleanup map + format cache AFTER the remote delete + removeUploadedSegment(name); } } @@ -723,7 +740,7 @@ String getMetadataFileForCommit(long primaryTerm, long generation) throws IOExce private void postUpload(Directory from, String src, String remoteFilename, String checksum) throws IOException { UploadedSegmentMetadata segmentMetadata = new UploadedSegmentMetadata(src, remoteFilename, checksum, from.fileLength(src)); - segmentsUploadedToRemoteStore.put(src, segmentMetadata); + addUploadedSegment(src, segmentMetadata); } /** @@ -816,6 +833,80 @@ public void uploadMetadata( } } + /** + * Upload metadata file using CatalogSnapshot. + * Uses polymorphic dispatch to CatalogSnapshot subclasses for Lucene version resolution + * and serialization, eliminating instanceof checks. + * + * @param segmentFiles segment files that are part of the shard at the time of the latest refresh + * @param catalogSnapshot CatalogSnapshot containing segment metadata (either SegmentInfos-backed or Composite) + * @param storeDirectory instance of local directory to temporarily create metadata file before upload + * @param translogGeneration translog generation + * @param replicationCheckpoint ReplicationCheckpoint of primary shard + * @param nodeId node id + * @throws IOException in case of I/O error while uploading the metadata file + */ + public void uploadMetadata( + Collection segmentFiles, + CatalogSnapshot catalogSnapshot, + Directory storeDirectory, + long translogGeneration, + ReplicationCheckpoint replicationCheckpoint, + String nodeId + ) throws IOException { + synchronized (this) { + String metadataFilename = MetadataFilenameUtils.getMetadataFilename( + replicationCheckpoint.getPrimaryTerm(), + catalogSnapshot.getGeneration(), + translogGeneration, + metadataUploadCounter.incrementAndGet(), + RemoteSegmentMetadata.CURRENT_VERSION, + nodeId + ); + try { + try (IndexOutput indexOutput = storeDirectory.createOutput(metadataFilename, IOContext.DEFAULT)) { + Map uploadedSegments = new HashMap<>(); + + // Polymorphic dispatch — no instanceof checks needed. + // Each CatalogSnapshot subclass knows how to resolve Lucene versions for its files. + for (String file : segmentFiles) { + if (segmentsUploadedToRemoteStore.containsKey(file)) { + UploadedSegmentMetadata metadata = segmentsUploadedToRemoteStore.get(file); + metadata.setWrittenByMajor(catalogSnapshot.getFormatVersionForFile(metadata.originalFilename)); + uploadedSegments.put(file, metadata.toString()); + } else { + throw new NoSuchFileException(file); + } + } + + // Polymorphic dispatch — each CatalogSnapshot subclass knows how to serialize itself + // to SegmentInfos bytes for the remote metadata file. + byte[] segmentInfoSnapshotByteArray = catalogSnapshot.serialize(); + + metadataStreamWrapper.writeStream( + indexOutput, + new RemoteSegmentMetadata( + RemoteSegmentMetadata.fromMapOfStrings(uploadedSegments), + segmentInfoSnapshotByteArray, + replicationCheckpoint + ) + ); + } + storeDirectory.sync(Collections.singleton(metadataFilename)); + remoteMetadataDirectory.copyFrom(storeDirectory, metadataFilename, metadataFilename, IOContext.DEFAULT); + } finally { + tryAndDeleteLocalFile(metadataFilename, storeDirectory); + } + } + } + + // TODO: When RemoteStoreRefreshListener is migrated to use CatalogSnapshot-based uploadMetadata, + // the instanceof check for SegmentInfosCatalogSnapshot.setUserData() will no longer be needed + // since setUserData() is now properly implemented in SegmentInfosCatalogSnapshot. + // Also, the old uploadMetadata(SegmentInfos, ...) overload above can be removed at that point + // and getSegmentToLuceneVersion() can be deleted since it's encapsulated in + // SegmentInfosCatalogSnapshot.getFormatVersionForFile(). + /** * Parses the provided SegmentInfos to retrieve a mapping of the provided segment files to * the respective Lucene major version that wrote the segments @@ -865,7 +956,24 @@ private void tryAndDeleteLocalFile(String filename, Directory directory) { } } + /** + * Gets the checksum of a local file by delegating to the local directory's + * format-aware checksum calculation. + * + *

    + *
  • If the local directory is a {@link DataFormatAwareStoreDirectory}, it uses + * {@code calculateUploadChecksum()} which routes via the registry.
    • + *
  • Otherwise, falls back to Lucene's {@code CodecUtil.retrieveChecksum()} for + * backward compatibility with non-composite directories.
    • + *
+ * + */ private String getChecksumOfLocalFile(Directory directory, String file) throws IOException { + DataFormatAwareStoreDirectory dfasd = DataFormatAwareStoreDirectory.unwrap(directory); + if (dfasd != null) { + return dfasd.calculateUploadChecksum(file); + } + // Fallback for non-optimized indices (backward compatibility) try (IndexInput indexInput = directory.openInput(file, IOContext.READONCE)) { return Long.toString(CodecUtil.retrieveChecksum(indexInput)); } @@ -880,8 +988,88 @@ public String getExistingRemoteFilename(String localFilename) { return null; } + // ═══════════════════════════════════════════════════════════════ + // Encapsulated map mutation APIs — keep format cache in sync + // IMPORTANT: All mutations to segmentsUploadedToRemoteStore and + // pendingDownloadMergedSegments must go through these methods. + // ═══════════════════════════════════════════════════════════════ + + /** Extract format from originalFilename. Returns "lucene" if no format prefix present. */ + private static String extractFormat(String originalFilename) { + return FileMetadata.parseDataFormat(originalFilename); + } + + /** Shared helper for single format registration. */ + private void registerFormatForBlob(String blobKey, String originalFilename) { + if (formatBlobRouter != null) { + formatBlobRouter.registerBlobFormat(blobKey, extractFormat(originalFilename)); + } + } + + /** Shared helper for single format unregistration. */ + private void unregisterFormatForBlob(String blobKey) { + if (formatBlobRouter != null) { + formatBlobRouter.unregisterBlobFormat(blobKey); + } + } + + /** + * Replace entire uploaded segments map + rebuild format cache. + * Called by init(), initializeToSpecificCommit(), initializeToSpecificTimestamp(). + */ + private void replaceUploadedSegments(Map newSegments) { + this.segmentsUploadedToRemoteStore = new ConcurrentHashMap<>(newSegments); + syncBlobFormatCache(); + } + + /** + * Add a segment to uploaded map + register its format in cache. + * Called by postUpload(). + */ + private void addUploadedSegment(String localFilename, UploadedSegmentMetadata metadata) { + segmentsUploadedToRemoteStore.put(localFilename, metadata); + registerFormatForBlob(metadata.getUploadedFilename(), metadata.getOriginalFilename()); + } + + /** + * Remove a segment from uploaded map + unregister its format from cache. + * Called by deleteFile(), deleteStaleSegments(). + * IMPORTANT: Call this AFTER the remote delete operation, not before, + * so that the format cache entry is still available during routing. + */ + private void removeUploadedSegment(String localFilename) { + UploadedSegmentMetadata removed = segmentsUploadedToRemoteStore.remove(localFilename); + if (removed != null) { + unregisterFormatForBlob(removed.getUploadedFilename()); + } + } + + /** + * Rebuild format cache from both segmentsUploadedToRemoteStore and pendingDownloadMergedSegments. + * Called after bulk replacement of segmentsUploadedToRemoteStore (init paths). + */ + private void syncBlobFormatCache() { + if (formatBlobRouter == null) { + return; + } + Map blobKeyToFormat = new HashMap<>(); + for (UploadedSegmentMetadata metadata : segmentsUploadedToRemoteStore.values()) { + blobKeyToFormat.put(metadata.getUploadedFilename(), extractFormat(metadata.getOriginalFilename())); + } + if (pendingDownloadMergedSegments != null) { + for (Map.Entry entry : pendingDownloadMergedSegments.entrySet()) { + blobKeyToFormat.put(entry.getValue(), extractFormat(entry.getKey())); + } + } + formatBlobRouter.replaceBlobFormatCache(blobKeyToFormat); + } + private String getNewRemoteSegmentFilename(String localFilename) { - return localFilename + SEGMENT_NAME_UUID_SEPARATOR + UUIDs.base64UUID(); + // Strip format prefix if present before appending UUID. + // For optimized indices, localFilename may be "format/filename" (e.g., "parquet/_0.pqt"). + // The blob key should be "filename__UUID" (e.g., "_0.pqt__UUID"), not "parquet/_0.pqt__UUID". + String plainFilename = FileMetadata.parseFile(localFilename); + return plainFilename + SEGMENT_NAME_UUID_SEPARATOR + UUIDs.base64UUID(); } private String getLocalSegmentFilename(String remoteFilename) { @@ -1059,7 +1247,7 @@ public void deleteStaleSegments(int lastNMetadataFilesToKeep) throws IOException // Update cache after successful batch deletion for (String file : filesToDelete) { if (!activeSegmentFilesMetadataMap.containsKey(getLocalSegmentFilename(file))) { - segmentsUploadedToRemoteStore.remove(getLocalSegmentFilename(file)); + removeUploadedSegment(getLocalSegmentFilename(file)); } } } catch (IOException e) { @@ -1182,6 +1370,10 @@ public void close() throws IOException { */ public void markMergedSegmentsPendingDownload(Map localToRemoteFilenames) { pendingDownloadMergedSegments.putAll(localToRemoteFilenames); + // Register format for each pending segment in the format cache + for (Map.Entry entry : localToRemoteFilenames.entrySet()) { + registerFormatForBlob(entry.getValue(), entry.getKey()); + } } /** @@ -1190,7 +1382,12 @@ public void markMergedSegmentsPendingDownload(Map localToRemoteF * @param localFilenames Set of local filenames to remove from pending downloads */ public void unmarkMergedSegmentsPendingDownload(Set localFilenames) { - localFilenames.forEach(pendingDownloadMergedSegments::remove); + for (String localFilename : localFilenames) { + String remoteFilename = pendingDownloadMergedSegments.remove(localFilename); + if (remoteFilename != null) { + unregisterFormatForBlob(remoteFilename); + } + } } /** diff --git a/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectoryFactory.java b/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectoryFactory.java index 53e8566ca8c90..5a2bb215f809c 100644 --- a/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectoryFactory.java +++ b/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectoryFactory.java @@ -8,17 +8,20 @@ package org.opensearch.index.store; +import org.apache.logging.log4j.LogManager; import org.apache.lucene.store.Directory; import org.apache.lucene.store.LockFactory; import org.opensearch.common.annotation.PublicApi; import org.opensearch.common.blobstore.BlobPath; import org.opensearch.core.index.shard.ShardId; import org.opensearch.index.IndexSettings; +import org.opensearch.index.engine.dataformat.DataFormatRegistry; import org.opensearch.index.remote.RemoteStorePathStrategy; import org.opensearch.index.remote.RemoteStoreUtils; import org.opensearch.index.shard.ShardPath; import org.opensearch.index.store.lockmanager.RemoteStoreLockManager; import org.opensearch.index.store.lockmanager.RemoteStoreLockManagerFactory; +import org.opensearch.index.store.remote.DataFormatAwareRemoteDirectory; import org.opensearch.plugins.IndexStorePlugin; import org.opensearch.repositories.RepositoriesService; import org.opensearch.repositories.Repository; @@ -48,15 +51,26 @@ public class RemoteSegmentStoreDirectoryFactory implements IndexStorePlugin.Dire private final String segmentsPathFixedPrefix; private final ThreadPool threadPool; + private final DataFormatRegistry dataFormatRegistry; public RemoteSegmentStoreDirectoryFactory( Supplier repositoriesService, ThreadPool threadPool, String segmentsPathFixedPrefix + ) { + this(repositoriesService, threadPool, segmentsPathFixedPrefix, null); + } + + public RemoteSegmentStoreDirectoryFactory( + Supplier repositoriesService, + ThreadPool threadPool, + String segmentsPathFixedPrefix, + DataFormatRegistry dataFormatRegistry ) { this.repositoriesService = repositoriesService; this.segmentsPathFixedPrefix = segmentsPathFixedPrefix; this.threadPool = threadPool; + this.dataFormatRegistry = dataFormatRegistry; } @Override @@ -75,7 +89,8 @@ public Directory newDirectory(IndexSettings indexSettings, ShardPath path) throw indexSettings.getRemoteStorePathStrategy(), null, RemoteStoreUtils.isServerSideEncryptionEnabledIndex(indexSettings.getIndexMetadata()), - indexSettings.isWarmIndex() + indexSettings.isWarmIndex(), + indexSettings ); } @@ -113,6 +128,28 @@ public Directory newDirectory( String indexFixedPrefix, boolean isServerSideEncryptionEnabled, boolean isWarmIndex + ) throws IOException { + return newDirectory( + repositoryName, + indexUUID, + shardId, + pathStrategy, + indexFixedPrefix, + isServerSideEncryptionEnabled, + isWarmIndex, + null + ); + } + + public Directory newDirectory( + String repositoryName, + String indexUUID, + ShardId shardId, + RemoteStorePathStrategy pathStrategy, + String indexFixedPrefix, + boolean isServerSideEncryptionEnabled, + boolean isWarmIndex, + IndexSettings indexSettings ) throws IOException { assert Objects.nonNull(pathStrategy); // We should be not calling close for repository. @@ -134,18 +171,30 @@ public Directory newDirectory( .indexFixedPrefix(indexFixedPrefix) .build(); - // Derive the path for data directory of SEGMENTS BlobPath dataPath = pathStrategy.generatePath(dataPathInput); - RemoteDirectory dataDirectory = new RemoteDirectory( - blobStoreRepository.blobStore(isServerSideEncryptionEnabled).blobContainer(dataPath), - blobStoreRepository::maybeRateLimitRemoteUploadTransfers, - blobStoreRepository::maybeRateLimitLowPriorityRemoteUploadTransfers, - isWarmIndex - ? blobStoreRepository::maybeRateLimitRemoteDownloadTransfersForWarm - : blobStoreRepository::maybeRateLimitRemoteDownloadTransfers, - blobStoreRepository::maybeRateLimitLowPriorityDownloadTransfers, - pendingDownloadMergedSegments - ); + RemoteDirectory dataDirectory = indexSettings != null && indexSettings.isPluggableDataFormatEnabled() + ? new DataFormatAwareRemoteDirectory( + blobStoreRepository.blobStore(isServerSideEncryptionEnabled), + dataPath, + blobStoreRepository::maybeRateLimitRemoteUploadTransfers, + blobStoreRepository::maybeRateLimitLowPriorityRemoteUploadTransfers, + blobStoreRepository::maybeRateLimitRemoteDownloadTransfers, + blobStoreRepository::maybeRateLimitLowPriorityDownloadTransfers, + pendingDownloadMergedSegments, + LogManager.getLogger("index.store.remote.composite." + shardId), + dataFormatRegistry, + indexSettings + ) + : new RemoteDirectory( + blobStoreRepository.blobStore(isServerSideEncryptionEnabled).blobContainer(dataPath), + blobStoreRepository::maybeRateLimitRemoteUploadTransfers, + blobStoreRepository::maybeRateLimitLowPriorityRemoteUploadTransfers, + isWarmIndex + ? blobStoreRepository::maybeRateLimitRemoteDownloadTransfersForWarm + : blobStoreRepository::maybeRateLimitRemoteDownloadTransfers, + blobStoreRepository::maybeRateLimitLowPriorityDownloadTransfers, + pendingDownloadMergedSegments + ); RemoteStorePathStrategy.ShardDataPathInput mdPathInput = RemoteStorePathStrategy.ShardDataPathInput.builder() .basePath(repositoryBasePath) diff --git a/server/src/main/java/org/opensearch/index/store/Store.java b/server/src/main/java/org/opensearch/index/store/Store.java index 3fadd56787ca2..624523ca6e24c 100644 --- a/server/src/main/java/org/opensearch/index/store/Store.java +++ b/server/src/main/java/org/opensearch/index/store/Store.java @@ -92,6 +92,8 @@ import org.opensearch.index.IndexSettings; import org.opensearch.index.engine.CombinedDeletionPolicy; import org.opensearch.index.engine.Engine; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; +import org.opensearch.index.engine.exec.coord.SegmentInfosCatalogSnapshot; import org.opensearch.index.seqno.SequenceNumbers; import org.opensearch.index.shard.AbstractIndexShardComponent; import org.opensearch.index.shard.IndexShard; @@ -231,6 +233,7 @@ public Store( this.isIndexSortEnabled = indexSettings.getIndexSortConfig().hasIndexSort(); this.isParentFieldEnabledVersion = indexSettings.getIndexVersionCreated().onOrAfter(org.opensearch.Version.V_3_2_0); this.directoryFactory = directoryFactory; + assert onClose != null; assert shardLock != null; assert shardLock.getShardId().equals(shardId); @@ -398,6 +401,31 @@ public Map getSegmentMetadataMap(SegmentInfos segment } } + /** + * Segment Replication method - Fetch a map of StoreFileMetadata for segments from a {@link CatalogSnapshot}, + * ignoring Segment_N files. Dispatches to the appropriate metadata loading strategy based on the snapshot type. + * + * @param catalogSnapshot {@link CatalogSnapshot} from which to compute metadata. + * @return {@link Map} map file name to {@link StoreFileMetadata}. + * @throws IOException in case of I/O error during metadata computation. + */ + // TODO: Remove the SegmentInfosCatalogSnapshot instanceof check once loadMetadata(CatalogSnapshot, ...) is fully implemented. + public Map getSegmentMetadataMap(CatalogSnapshot catalogSnapshot) throws IOException { + assert indexSettings.isSegRepEnabledOrRemoteNode(); + failIfCorrupted(); + + if (catalogSnapshot instanceof SegmentInfosCatalogSnapshot segmentInfosCatalogSnapshot) { + return getSegmentMetadataMap(segmentInfosCatalogSnapshot.getSegmentInfos()); + } + + try { + return loadMetadata(catalogSnapshot, directory, logger, true).fileMetadata; + } catch (NoSuchFileException | CorruptIndexException | IndexFormatTooOldException | IndexFormatTooNewException ex) { + markStoreCorrupted(ex); + throw ex; + } + } + /** * Segment Replication method * Returns a diff between the Maps of StoreFileMetadata that can be used for getting list of files to copy over to a replica for segment replication. The returned diff will hold a list of files that are: @@ -571,7 +599,7 @@ private void closeInternal() { // Leverage try-with-resources to close the shard lock for us try (Closeable c = shardLock) { try { - directory.innerClose(); // this closes the distributorDirectory as well + directory.innerClose(); // this closes the entire directory chain including DataFormatAwareStoreDirectory } finally { onClose.accept(shardLock); } @@ -1212,6 +1240,16 @@ public static LoadedMetadata loadMetadata(SegmentInfos segmentInfos, Directory d return new LoadedMetadata(unmodifiableMap(builder), unmodifiableMap(commitUserDataBuilder), numDocs); } + public static LoadedMetadata loadMetadata( + CatalogSnapshot catalogSnapshot, + Directory directory, + Logger logger, + boolean ignoreSegmentsFile + ) throws IOException { + // TODO: Implement format-aware loadMetadata equivalent to the SegmentInfos version + throw new UnsupportedOperationException("loadMetadata for CatalogSnapshot is not yet implemented"); + } + private static void checksumFromLuceneFile( Directory directory, String file, diff --git a/server/src/main/java/org/opensearch/index/store/SubdirectoryAwareDirectory.java b/server/src/main/java/org/opensearch/index/store/SubdirectoryAwareDirectory.java new file mode 100644 index 0000000000000..980c82b490e66 --- /dev/null +++ b/server/src/main/java/org/opensearch/index/store/SubdirectoryAwareDirectory.java @@ -0,0 +1,140 @@ +/* + * 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.store; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FilterDirectory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.IndexOutput; +import org.opensearch.index.shard.ShardPath; + +import java.io.IOException; +import java.nio.file.FileVisitResult; +import java.nio.file.Files; +import java.nio.file.NoSuchFileException; +import java.nio.file.Path; +import java.nio.file.SimpleFileVisitor; +import java.nio.file.attribute.BasicFileAttributes; +import java.util.Arrays; +import java.util.Collection; +import java.util.HashSet; +import java.util.Set; +import java.util.stream.Collectors; + +/** + * A Lucene Directory implementation that handles files in subdirectories. + * + * This directory wrapper enables file operations across subdirectories within + * the shard data path. It resolves paths, creates necessary directory structures, + * and delegates actual file operations to appropriate filesystem locations. + */ +public class SubdirectoryAwareDirectory extends FilterDirectory { + private static final Logger logger = LogManager.getLogger(SubdirectoryAwareDirectory.class); + private static final Set EXCLUDED_SUBDIRECTORIES = Set.of("index/", "translog/", "_state/"); + private final ShardPath shardPath; + + /** + * Constructor for SubdirectoryAwareDirectory. + * + * @param delegate the delegate directory + * @param shardPath the shard path + */ + public SubdirectoryAwareDirectory(Directory delegate, ShardPath shardPath) { + super(delegate); + this.shardPath = shardPath; + } + + @Override + public IndexInput openInput(String name, IOContext context) throws IOException { + return super.openInput(parseFilePath(name), context); + } + + @Override + public IndexOutput createOutput(String name, IOContext context) throws IOException { + String targetFilePath = parseFilePath(name); + Path targetFile = Path.of(targetFilePath); + Files.createDirectories(targetFile.getParent()); + return super.createOutput(targetFilePath, context); + } + + @Override + public void deleteFile(String name) throws IOException { + super.deleteFile(parseFilePath(name)); + } + + @Override + public long fileLength(String name) throws IOException { + return super.fileLength(parseFilePath(name)); + } + + @Override + public void sync(Collection names) throws IOException { + super.sync(names.stream().map(this::parseFilePath).collect(Collectors.toList())); + } + + @Override + public void rename(String source, String dest) throws IOException { + super.rename(parseFilePath(source), parseFilePath(dest)); + } + + @Override + public String[] listAll() throws IOException { + // Get files from the delegate (regular index files) + String[] delegateFiles = super.listAll(); + + // Get subdirectory files by scanning all subdirectories + Set allFiles = new HashSet<>(Arrays.asList(delegateFiles)); + addSubdirectoryFiles(allFiles); + + return allFiles.stream().sorted().toArray(String[]::new); + } + + private void addSubdirectoryFiles(Set allFiles) throws IOException { + Path dataPath = shardPath.getDataPath(); + Files.walkFileTree(dataPath, new SimpleFileVisitor<>() { + @Override + public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) { + if (attrs.isRegularFile()) { + Path relativePath = dataPath.relativize(file); + // Only add files that are in subdirectories (have a parent directory) + if (relativePath.getParent() != null) { + String relativePathStr = relativePath.toString(); + // Exclude index dir (handled in super.listAll()), translog dir, and _state dir + if (EXCLUDED_SUBDIRECTORIES.stream().noneMatch(relativePathStr::startsWith)) { + allFiles.add(relativePathStr); + } + } + } + return FileVisitResult.CONTINUE; + } + + @Override + public FileVisitResult visitFileFailed(Path file, IOException e) throws IOException { + if (e instanceof NoSuchFileException) { + logger.debug("Skipping inaccessible file during size estimation: {}", file); + return FileVisitResult.CONTINUE; + } + throw e; + } + }); + } + + private String parseFilePath(String fileName) { + if (Path.of(fileName).getParent() != null) { + // File path (e.g., "subdirectory/segments_1" or "subdirectory/recovery.xxx.segments_1") + return shardPath.getDataPath().resolve(fileName).toString(); + } else { + // Simple filename (e.g., "segments_1") - resolve relative to the shard's index directory + return shardPath.resolveIndex().resolve(fileName).toString(); + } + } +} diff --git a/server/src/main/java/org/opensearch/index/store/checksum/GenericCRC32ChecksumHandler.java b/server/src/main/java/org/opensearch/index/store/checksum/GenericCRC32ChecksumHandler.java new file mode 100644 index 0000000000000..cb69b48d410c0 --- /dev/null +++ b/server/src/main/java/org/opensearch/index/store/checksum/GenericCRC32ChecksumHandler.java @@ -0,0 +1,48 @@ +/* + * 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.store.checksum; + +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.opensearch.common.annotation.PublicApi; +import org.opensearch.index.store.FormatChecksumStrategy; + +import java.io.IOException; +import java.util.zip.CRC32; + +/** + * Checksum strategy that computes CRC32 over the entire file contents. + * + *

This is the default/fallback strategy for non-Lucene formats (Parquet, Arrow, etc.) + * that do not embed a Lucene codec footer. It reads the entire file — O(n) complexity.

+ * + * @opensearch.api + */ +@PublicApi(since = "3.0.0") +public class GenericCRC32ChecksumHandler implements FormatChecksumStrategy { + + private static final int BUFFER_SIZE = 8192; + + @Override + public long computeChecksum(Directory dir, String fileName) throws IOException { + CRC32 crc32 = new CRC32(); + byte[] buffer = new byte[BUFFER_SIZE]; + try (IndexInput input = dir.openInput(fileName, IOContext.READONCE)) { + long remaining = input.length(); + while (remaining > 0) { + int toRead = (int) Math.min(buffer.length, remaining); + input.readBytes(buffer, 0, toRead); + crc32.update(buffer, 0, toRead); + remaining -= toRead; + } + } + return crc32.getValue(); + } +} diff --git a/server/src/main/java/org/opensearch/index/store/checksum/LuceneChecksumHandler.java b/server/src/main/java/org/opensearch/index/store/checksum/LuceneChecksumHandler.java new file mode 100644 index 0000000000000..3776cb8ff4285 --- /dev/null +++ b/server/src/main/java/org/opensearch/index/store/checksum/LuceneChecksumHandler.java @@ -0,0 +1,37 @@ +/* + * 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.store.checksum; + +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.opensearch.common.annotation.PublicApi; +import org.opensearch.index.store.FormatChecksumStrategy; + +import java.io.IOException; + +/** + * Checksum strategy for Lucene segment files. + * + *

Reads the checksum from the Lucene codec footer — an O(1) operation + * since it only reads the last 16 bytes of the file.

+ * + * @opensearch.api + */ +@PublicApi(since = "3.0.0") +public class LuceneChecksumHandler implements FormatChecksumStrategy { + + @Override + public long computeChecksum(Directory dir, String fileName) throws IOException { + try (IndexInput input = dir.openInput(fileName, IOContext.READONCE)) { + return CodecUtil.retrieveChecksum(input); + } + } +} diff --git a/server/src/main/java/org/opensearch/index/store/checksum/package-info.java b/server/src/main/java/org/opensearch/index/store/checksum/package-info.java new file mode 100644 index 0000000000000..1b27fcadbfc63 --- /dev/null +++ b/server/src/main/java/org/opensearch/index/store/checksum/package-info.java @@ -0,0 +1,14 @@ +/* + * 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. + */ + +/** + * Checksum calculation handlers for format-aware file integrity verification. + * Provides a registry-based approach to delegate checksum computation to format-specific + * handlers (e.g., Lucene CodecUtil for Lucene files, CRC32 for other formats). + */ +package org.opensearch.index.store.checksum; diff --git a/server/src/main/java/org/opensearch/index/store/remote/DataFormatAwareRemoteDirectory.java b/server/src/main/java/org/opensearch/index/store/remote/DataFormatAwareRemoteDirectory.java new file mode 100644 index 0000000000000..6b17f1f66ae23 --- /dev/null +++ b/server/src/main/java/org/opensearch/index/store/remote/DataFormatAwareRemoteDirectory.java @@ -0,0 +1,565 @@ +/* + * 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.store.remote; + +import org.apache.logging.log4j.Logger; +import org.apache.logging.log4j.message.ParameterizedMessage; +import org.apache.lucene.index.CorruptIndexException; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.IndexOutput; +import org.opensearch.ExceptionsHelper; +import org.opensearch.cluster.metadata.CryptoMetadata; +import org.opensearch.common.annotation.InternalApi; +import org.opensearch.common.blobstore.AsyncMultiStreamBlobContainer; +import org.opensearch.common.blobstore.BlobContainer; +import org.opensearch.common.blobstore.BlobMetadata; +import org.opensearch.common.blobstore.BlobPath; +import org.opensearch.common.blobstore.BlobStore; +import org.opensearch.common.blobstore.exception.CorruptFileException; +import org.opensearch.common.blobstore.stream.write.WriteContext; +import org.opensearch.common.blobstore.stream.write.WritePriority; +import org.opensearch.common.blobstore.transfer.RemoteTransferContainer; +import org.opensearch.common.blobstore.transfer.stream.OffsetRangeIndexInputStream; +import org.opensearch.common.blobstore.transfer.stream.OffsetRangeInputStream; +import org.opensearch.common.lucene.store.ByteArrayIndexInput; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.common.unit.ByteSizeUnit; +import org.opensearch.index.IndexSettings; +import org.opensearch.index.engine.dataformat.DataFormatRegistry; +import org.opensearch.index.store.DataFormatAwareStoreDirectory; +import org.opensearch.index.store.FileMetadata; +import org.opensearch.index.store.RemoteDirectory; +import org.opensearch.index.store.RemoteIndexInput; +import org.opensearch.index.store.RemoteIndexOutput; + +import java.io.IOException; +import java.io.InputStream; +import java.nio.file.NoSuchFileException; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.function.UnaryOperator; + +/** + * DataFormatAwareRemoteDirectory extends RemoteDirectory with format-aware blob routing. + * + *

This directory routes file operations to format-specific BlobContainers. Format resolution + * depends on the caller: + *

    + *
  • Remote blob operations (deleteFile, openInput, fileLength, openBlockInput): receive plain + * blob keys (e.g., "_0.pqt__UUID") from RSSD and resolve format via {@code blobFormatCache}
    • + *
  • Upload operations (copyFrom): receive local filenames with "format/file" convention + * (e.g., "parquet/_0.pqt") and parse format directly via {@link FileMetadata}
    • + *
  • FileMetadata-based APIs: receive format explicitly via the FileMetadata object
    • + *
+ * + *

Blob container routing: + *

    + *
  • "lucene" format (or no format) → inherited blobContainer at baseBlobPath (same as RemoteDirectory)
    • + *
  • Non-lucene formats (e.g., "parquet") → baseBlobPath/formatName/ sub-path
    • + *
+ * + * @opensearch.api + */ +@InternalApi +public class DataFormatAwareRemoteDirectory extends RemoteDirectory { + + private static final String DEFAULT_FORMAT = "lucene"; + + private final UnaryOperator uploadRateLimiter; + private final UnaryOperator lowPriorityUploadRateLimiter; + private final DownloadRateLimiterProvider downloadRateLimiterProvider; + + private final FormatBlobRouter formatBlobRouter; + private final Logger logger; + + /** + * Full constructor with all rate limiter parameters. + */ + public DataFormatAwareRemoteDirectory( + BlobStore blobStore, + BlobPath baseBlobPath, + UnaryOperator uploadRateLimiter, + UnaryOperator lowPriorityUploadRateLimiter, + UnaryOperator downloadRateLimiter, + UnaryOperator lowPriorityDownloadRateLimiter, + Map pendingDownloadMergedSegments, + Logger logger, + DataFormatRegistry dataFormatRegistry, + IndexSettings indexSettings + ) { + super( + blobStore.blobContainer(baseBlobPath), + uploadRateLimiter, + lowPriorityUploadRateLimiter, + downloadRateLimiter, + lowPriorityDownloadRateLimiter, + pendingDownloadMergedSegments + ); + this.formatBlobRouter = new FormatBlobRouter(blobStore, baseBlobPath); + this.uploadRateLimiter = uploadRateLimiter; + this.lowPriorityUploadRateLimiter = lowPriorityUploadRateLimiter; + this.downloadRateLimiterProvider = new DownloadRateLimiterProvider(downloadRateLimiter, lowPriorityDownloadRateLimiter); + this.logger = logger; + + // Pre-register format-specific BlobContainers from DataFormatRegistry + if (dataFormatRegistry != null && indexSettings != null) { + for (String formatName : dataFormatRegistry.getFormatDescriptors(indexSettings).keySet()) { + formatBlobRouter.registerFormat(formatName); + } + } + + logger.debug("Created DataFormatAwareRemoteDirectory with formats: {}", formatBlobRouter.registeredFormats()); + } + + // ═══════════════════════════════════════════════════════════════ + // Format Routing — delegates to FormatBlobRouter + // ═══════════════════════════════════════════════════════════════ + + /** + * Resolve the data format for a plain blob key using the format cache. + * Delegates to {@link FormatBlobRouter#resolveFormat(String)}. + * + * @param name the blob key to resolve format for + * @return the resolved data format name, defaults to "lucene" + */ + private String resolveFormat(String name) { + return formatBlobRouter.resolveFormat(name); + } + + /** + * Get BlobContainer for a specific data format. + * Delegates to {@link FormatBlobRouter#containerFor(String)}. + * + * @param format the data format name (e.g., "lucene", "parquet") + * @return BlobContainer for the format + */ + public BlobContainer getBlobContainerForFormat(String format) { + return formatBlobRouter.containerFor(format); + } + + /** + * Returns the {@link FormatBlobRouter} for direct access by callers that need + * format-aware blob operations (e.g., listing all blobs across formats). + * + * @return the format blob router + */ + @Override + public Optional getFormatBlobRouter() { + return Optional.of(formatBlobRouter); + } + + // ═══════════════════════════════════════════════════════════════ + // Aggregated operations across all format containers + // ═══════════════════════════════════════════════════════════════ + + /** + * Lists all blobs across the base container and all format-specific containers. + * Results are sorted in UTF-16 order as required by the Directory contract. + */ + @Override + public String[] listAll() throws IOException { + Map allBlobs = formatBlobRouter.listAllBlobs(); + String[] result = allBlobs.keySet().toArray(new String[0]); + Arrays.sort(result); + return result; + } + + /** + * Format-aware deleteFile override. + * + *

Uses {@link #resolveFormat(String)} to determine which BlobContainer to delete from. + * Name is always a plain blob key (e.g., "_0.pqt__UUID") from RSSD, resolved via blobFormatCache. + * Falls back to "lucene" (base container) if no format info is available. + */ + @Override + public void deleteFile(String name) throws IOException { + String format = resolveFormat(name); + BlobContainer container = getBlobContainerForFormat(format); + container.deleteBlobsIgnoringIfNotExists(Collections.singletonList(name)); + } + + /** + * Format-aware batch delete override. + * + *

Note: This method receives plain blob names (e.g., "_0.parquet__UUID") without format info, + * so it attempts deletion from ALL containers (base + format-specific). This is used during + * stale segment cleanup where the caller doesn't have format information. + */ + @Override + public void deleteFiles(List names) throws IOException { + if (names == null || names.isEmpty()) { + return; + } + // Delete from base container (handles lucene/metadata blobs) + super.deleteFiles(names); + + // Broadcast delete to every format-specific container. This is intentionally speculative: + // blob names are UUID-suffixed and globally unique, so at most one container holds each + // blob. + for (String format : formatBlobRouter.registeredFormats()) { + formatBlobRouter.containerFor(format).deleteBlobsIgnoringIfNotExists(names); + } + } + + // ═══════════════════════════════════════════════════════════════ + // String-based overrides — called by RemoteSegmentStoreDirectory + // These parse "format/file" from the src string + // ═══════════════════════════════════════════════════════════════ + + /** + * Sync copyFrom override that properly handles format-aware local files. + * + *

When AsyncMultiStreamBlobContainer is not available (e.g., FS-based blob store in tests), + * this fallback is used. The src string may contain "format/" prefix (e.g., "parquet/_0.pqt") + * which the source DataFormatAwareStoreDirectory handles via parseFilePath(). + */ + @Override + public void copyFrom(Directory from, String src, String dest, IOContext context) throws IOException { + logger.debug("Sync copyFrom: src={}, dest={}", src, dest); + FileMetadata fileMetadata = new FileMetadata(src); + BlobContainer container = getBlobContainerForFormat(fileMetadata.dataFormat()); + // Read from local directory (DataFormatAwareStoreDirectory handles "format/file" in src) + // Write to format-specific BlobContainer (lucene→base, parquet→parquet sub-path, etc.) + try (IndexInput is = from.openInput(src, context); IndexOutput os = new RemoteIndexOutput(dest, container)) { + os.copyBytes(is, is.length()); + } + } + + /** + * Format-aware async copyFrom override. + * + *

Parses the src string (e.g., "parquet/_0.pqt") to determine format routing. + * Opens local file using src as-is (DataFormatAwareStoreDirectory handles format/file parsing). + * Uploads to the format-specific BlobContainer using remoteFileName as the blob key. + */ + @Override + public boolean copyFrom( + Directory from, + String src, + String remoteFileName, + IOContext context, + Runnable postUploadRunner, + ActionListener listener, + boolean lowPriorityUpload, + CryptoMetadata cryptoMetadata + ) { + try { + FileMetadata fileMetadata = new FileMetadata(src); + BlobContainer container = getBlobContainerForFormat(fileMetadata.dataFormat()); + + if (container instanceof AsyncMultiStreamBlobContainer) { + logger.debug( + "Format-aware upload: src={}, format={}, remoteFile={}, container={}", + src, + fileMetadata.dataFormat(), + remoteFileName, + container.path() + ); + uploadBlob(from, src, remoteFileName, container, context, postUploadRunner, listener, lowPriorityUpload, cryptoMetadata); + return true; + } + + logger.warn("BlobContainer for format {} does not support async multi-stream upload", fileMetadata.dataFormat()); + return false; + } catch (Exception e) { + logger.error(() -> new ParameterizedMessage("Failed format-aware upload: src={}, error={}", src, e.getMessage()), e); + listener.onFailure(e); + return true; // Handled (even though failed) + } + } + + /** + * Format-aware fileLength override. + * + *

Receives a plain blob key (e.g., "_0.pqt__UUID") from RSSD and uses + * {@link #resolveFormat(String)} to look up the format from blobFormatCache. + */ + @Override + public long fileLength(String name) throws IOException { + String format = resolveFormat(name); + BlobContainer container = getBlobContainerForFormat(format); + + if (container == null) { + throw new NoSuchFileException(String.format(java.util.Locale.ROOT, "No container for format %s, file %s", format, name)); + } + + List metadata = container.listBlobsByPrefixInSortedOrder(name, 1, BlobContainer.BlobNameSortOrder.LEXICOGRAPHIC); + if (metadata.size() == 1 && metadata.get(0).name().equals(name)) { + return metadata.get(0).length(); + } + throw new NoSuchFileException(name); + } + + /** + * Create output for a specific format. + */ + public RemoteIndexOutput createOutput(String remoteFileName, String dataFormat, IOContext context) throws IOException { + BlobContainer container = getBlobContainerForFormat(dataFormat); + if (container == null) { + throw new IOException(String.format(java.util.Locale.ROOT, "No container for format %s, file %s", dataFormat, remoteFileName)); + } + return new RemoteIndexOutput(remoteFileName, container); + } + + // ═══════════════════════════════════════════════════════════════ + // Lifecycle + // ═══════════════════════════════════════════════════════════════ + + @Override + public void delete() throws IOException { + // Delete all format-specific containers + for (String format : formatBlobRouter.registeredFormats()) { + formatBlobRouter.containerFor(format).delete(); + } + // Also delete the base container (inherited from RemoteDirectory) + super.delete(); + logger.debug("Deleted all containers from DataFormatAwareRemoteDirectory"); + } + + @Override + public void close() throws IOException { + formatBlobRouter.clearBlobFormatCache(); + } + + @Override + public String toString() { + return "DataFormatAwareRemoteDirectory{" + + "formats=" + + formatBlobRouter.registeredFormats() + + ", basePath=" + + formatBlobRouter.basePath() + + '}'; + } + + // ═══════════════════════════════════════════════════════════════ + // Private upload helpers + // ═══════════════════════════════════════════════════════════════ + + /** + * Upload blob using String-based src. Opens local file from 'from' directory using src as-is. + * The target BlobContainer is determined by the caller (format-aware routing already done). + */ + private void uploadBlob( + Directory from, + String src, + String remoteFileName, + BlobContainer targetContainer, + IOContext ioContext, + Runnable postUploadRunner, + ActionListener listener, + boolean lowPriorityUpload, + CryptoMetadata cryptoMetadata + ) throws Exception { + assert ioContext != IOContext.READONCE : "Remote upload will fail with IoContext.READONCE"; + long expectedChecksum; + DataFormatAwareStoreDirectory dfasd = DataFormatAwareStoreDirectory.unwrap(from); + if (dfasd != null) { + expectedChecksum = dfasd.calculateChecksum(src); + } else { + expectedChecksum = calculateChecksumOfChecksum(from, src); + } + IndexInput indexInput = from.openInput(src, ioContext); + try { + long contentLength = indexInput.length(); + boolean remoteIntegrityEnabled = (targetContainer instanceof AsyncMultiStreamBlobContainer) + && ((AsyncMultiStreamBlobContainer) targetContainer).remoteIntegrityCheckSupported(); + + lowPriorityUpload = lowPriorityUpload || contentLength > ByteSizeUnit.GB.toBytes(15); + + RemoteTransferContainer.OffsetRangeInputStreamSupplier supplier = lowPriorityUpload + ? (size, position) -> lowPriorityUploadRateLimiter.apply( + new OffsetRangeIndexInputStream(indexInput.clone(), size, position) + ) + : (size, position) -> uploadRateLimiter.apply(new OffsetRangeIndexInputStream(indexInput.clone(), size, position)); + + RemoteTransferContainer remoteTransferContainer = new RemoteTransferContainer( + src, + remoteFileName, + contentLength, + true, + lowPriorityUpload ? WritePriority.LOW : WritePriority.NORMAL, + supplier, + expectedChecksum, + remoteIntegrityEnabled + ); + + ActionListener completionListener = createCompletionListener( + src, + postUploadRunner, + listener, + remoteTransferContainer, + indexInput + ); + + WriteContext writeContext = remoteTransferContainer.createWriteContext(); + ((AsyncMultiStreamBlobContainer) targetContainer).asyncBlobUpload(writeContext, completionListener); + } catch (Exception e) { + logger.warn("Exception while calling asyncBlobUpload for {}, closing IndexInput", src); + indexInput.close(); + throw e; + } + } + + /** + * Opens a stream for reading the existing file and returns {@link RemoteIndexInput} enclosing + * the stream. + * + *

Receives a plain blob key (e.g., "_0.pqt__UUID") from RSSD and uses + * {@link #resolveFormat(String)} to look up the format from blobFormatCache. + * + * @param name the name of an existing file. + * @param fileLength file length + * @param context desired {@link IOContext} context + * @return the {@link RemoteIndexInput} enclosing the stream + * @throws IOException in case of I/O error + * @throws NoSuchFileException if the file does not exist + */ + @Override + public IndexInput openInput(String name, long fileLength, IOContext context) throws IOException { + String format = resolveFormat(name); + BlobContainer container = getBlobContainerForFormat(format); + InputStream inputStream = null; + try { + inputStream = container.readBlob(name); + UnaryOperator rateLimiter = downloadRateLimiterProvider.get(name); + return new RemoteIndexInput(name, rateLimiter.apply(inputStream), fileLength); + } catch (Exception e) { + // In case the RemoteIndexInput creation fails, close the input stream to avoid file handler leak. + if (inputStream != null) { + try { + inputStream.close(); + } catch (Exception closeEx) { + e.addSuppressed(closeEx); + } + } + logger.error("Exception while reading blob for file: {} format: {} path: {}", name, format, blobContainer.path()); + throw e; + } + } + + /** + * Format-aware openBlockInput override. + * + *

Receives a plain blob key (e.g., "_0.pqt__UUID") from RSSD and uses + * {@link #resolveFormat(String)} to look up the format from blobFormatCache. + * + * @param name the name of an existing file (blob key). + * @param position block start position + * @param length block length + * @param fileLength total file length + * @param context desired {@link IOContext} context + * @return the {@link IndexInput} enclosing the block data + * @throws IOException in case of I/O error + * @throws NoSuchFileException if the file does not exist + */ + @Override + public IndexInput openBlockInput(String name, long position, long length, long fileLength, IOContext context) throws IOException { + String format = resolveFormat(name); + BlobContainer container = getBlobContainerForFormat(format); + if (position < 0 || length <= 0 || (position + length > fileLength)) { + throw new IllegalArgumentException("Invalid values of block start and size"); + } + byte[] bytes; + try (InputStream inputStream = container.readBlob(name, position, length)) { + UnaryOperator rateLimiter = downloadRateLimiterProvider.get(name); + bytes = rateLimiter.apply(inputStream).readAllBytes(); + } catch (Exception e) { + logger.error("Exception while reading block for file: {} format: {} path: {}", name, format, blobContainer.path()); + throw e; + } + return new ByteArrayIndexInput(name, bytes); + } + + private ActionListener createCompletionListener( + String fileName, + Runnable postUploadRunner, + ActionListener listener, + RemoteTransferContainer remoteTransferContainer, + IndexInput indexInput + ) { + ActionListener completionListener = ActionListener.wrap(resp -> { + try { + postUploadRunner.run(); + listener.onResponse(null); + } catch (Exception e) { + logger.error(() -> new ParameterizedMessage("Exception in segment postUpload for file [{}]", fileName), e); + listener.onFailure(e); + } + }, ex -> { + logger.error(() -> new ParameterizedMessage("Failed to upload blob {}", fileName), ex); + IOException corruptIndexException = ExceptionsHelper.unwrapCorruption(ex); + if (corruptIndexException != null) { + listener.onFailure(corruptIndexException); + return; + } + Throwable throwable = ExceptionsHelper.unwrap(ex, CorruptFileException.class); + if (throwable != null) { + CorruptFileException cfe = (CorruptFileException) throwable; + listener.onFailure(new CorruptIndexException(cfe.getMessage(), cfe.getFileName())); + return; + } + listener.onFailure(ex); + }); + + completionListener = ActionListener.runBefore(completionListener, () -> { + try { + remoteTransferContainer.close(); + } catch (Exception e) { + logger.warn("Error closing RemoteTransferContainer", e); + } + }); + + completionListener = ActionListener.runAfter(completionListener, () -> { + try { + indexInput.close(); + } catch (IOException e) { + logger.warn("Error closing IndexInput", e); + } + }); + + return completionListener; + } + + // ═══════════════════════════════════════════════════════════════ + // Private helpers + // ═══════════════════════════════════════════════════════════════ + + private boolean isMergedSegment(String remoteFilename) { + return pendingDownloadMergedSegments != null && pendingDownloadMergedSegments.containsValue(remoteFilename); + } + + /** + * DownloadRateLimiterProvider returns a low-priority rate limited stream if the segment + * being downloaded is a merged segment. + */ + private class DownloadRateLimiterProvider { + private final UnaryOperator downloadRateLimiter; + private final UnaryOperator lowPriorityDownloadRateLimiter; + + DownloadRateLimiterProvider( + UnaryOperator downloadRateLimiter, + UnaryOperator lowPriorityDownloadRateLimiter + ) { + this.downloadRateLimiter = downloadRateLimiter; + this.lowPriorityDownloadRateLimiter = lowPriorityDownloadRateLimiter; + } + + public UnaryOperator get(final String filename) { + if (isMergedSegment(filename)) { + return lowPriorityDownloadRateLimiter; + } + return downloadRateLimiter; + } + } +} diff --git a/server/src/main/java/org/opensearch/index/store/remote/FormatBlobRouter.java b/server/src/main/java/org/opensearch/index/store/remote/FormatBlobRouter.java new file mode 100644 index 0000000000000..1e2e0922f3073 --- /dev/null +++ b/server/src/main/java/org/opensearch/index/store/remote/FormatBlobRouter.java @@ -0,0 +1,246 @@ +/* + * 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.store.remote; + +import org.opensearch.common.annotation.InternalApi; +import org.opensearch.common.blobstore.BlobContainer; +import org.opensearch.common.blobstore.BlobMetadata; +import org.opensearch.common.blobstore.BlobPath; +import org.opensearch.common.blobstore.BlobStore; + +import java.io.IOException; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.Locale; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; + +/** + * Routes blob operations to format-specific {@link BlobContainer}s. + * + *

Encapsulates the mapping from data format names (e.g., "lucene", "parquet") to + * their corresponding blob containers in the remote store. Formats that are considered + * "base path" formats (lucene, metadata) share the root blob container. All other + * formats get a sub-path container (e.g., {@code basePath/parquet/}). + * + *

Blob containers are created lazily on first access via {@link #containerFor(String)}, + * so new formats can be added without restarting the directory. + * + *

This class is thread-safe. Concurrent calls to {@link #containerFor(String)} for + * the same format will produce the same container instance (via {@code computeIfAbsent}). + * + * @opensearch.internal + */ +@InternalApi +public class FormatBlobRouter { + + /** Formats that route to the base blob container (same path as single-format RemoteDirectory). */ + private static final Set BASE_PATH_FORMATS = Set.of("lucene", "LUCENE", "metadata"); + + private static final String DEFAULT_FORMAT = "lucene"; + + private final BlobStore blobStore; + private final BlobPath basePath; + private final BlobContainer baseContainer; + private final ConcurrentHashMap formatContainers; + + /** + * Reverse lookup cache: maps remote blob keys (e.g., "_0.pqt__UUID") to their data format + * (e.g., "parquet"). Used by download/delete operations where the caller only has a plain + * blob key and needs to resolve which format container to route to. + * + *

Volatile reference to immutable map for atomic swap in {@link #replaceBlobFormatCache(Map)}. + */ + private volatile Map blobFormatCache = Map.of(); + + /** + * Creates a router with the given blob store and base path. + * The base container is created immediately; format-specific containers are lazy. + * + * @param blobStore the blob store for creating containers + * @param basePath the base blob path for this shard's remote segment store + */ + public FormatBlobRouter(BlobStore blobStore, BlobPath basePath) { + this.blobStore = blobStore; + this.basePath = basePath; + this.baseContainer = blobStore.blobContainer(basePath); + this.formatContainers = new ConcurrentHashMap<>(); + // Pre-register the default format so containerFor("lucene") doesn't create a sub-path + this.formatContainers.put(DEFAULT_FORMAT, baseContainer); + } + + /** + * Returns the blob container for the given format. + * + *

Base-path formats ("lucene", "metadata") return the root container. + * All other formats return a sub-path container at {@code basePath/formatName/}, + * created lazily on first access. + * + * @param format the data format name (e.g., "lucene", "parquet"). Null defaults to "lucene". + * @return the blob container for the format + */ + public BlobContainer containerFor(String format) { + if (format == null || format.isEmpty() || BASE_PATH_FORMATS.contains(format)) { + return baseContainer; + } + return formatContainers.computeIfAbsent(format, this::createFormatContainer); + } + + /** + * Returns the base blob container (used for lucene/metadata files). + * Equivalent to {@code containerFor("lucene")}. + * + * @return the base blob container + */ + public BlobContainer baseContainer() { + return baseContainer; + } + + /** + * Returns the base blob path. + * + * @return the base blob path + */ + public BlobPath basePath() { + return basePath; + } + + /** + * Lists all blobs across all known format containers. + * + *

Aggregates blobs from the base container and all format-specific containers + * that have been accessed (lazily created). The returned map is unmodifiable. + * + * @return map of blob name to metadata across all format containers + * @throws IOException if listing fails for any container + */ + public Map listAllBlobs() throws IOException { + Map all = new LinkedHashMap<>(baseContainer.listBlobs()); + for (Map.Entry entry : formatContainers.entrySet()) { + // Skip the default format — already listed via baseContainer + if (DEFAULT_FORMAT.equals(entry.getKey())) { + continue; + } + all.putAll(entry.getValue().listBlobs()); + } + return Collections.unmodifiableMap(all); + } + + /** + * Lists blobs in a specific format's container. + * + * @param format the data format name + * @return map of blob name to metadata for the format + * @throws IOException if listing fails + */ + public Map listBlobs(String format) throws IOException { + return containerFor(format).listBlobs(); + } + + /** + * Returns the set of all format names that have been accessed (have containers). + * Always includes "lucene". + * + * @return unmodifiable set of registered format names + */ + public Set registeredFormats() { + Set formats = new LinkedHashSet<>(); + formats.add(DEFAULT_FORMAT); + for (String format : formatContainers.keySet()) { + if (DEFAULT_FORMAT.equals(format) == false) { + formats.add(format); + } + } + return Collections.unmodifiableSet(formats); + } + + /** + * Pre-registers a format so its container is created eagerly. + * Useful during initialization when the set of formats is known from index settings. + * + * @param format the format name to pre-register + */ + public void registerFormat(String format) { + if (format != null && BASE_PATH_FORMATS.contains(format) == false) { + formatContainers.computeIfAbsent(format, this::createFormatContainer); + } + } + + // ═══════════════════════════════════════════════════════════════ + // Blob Format Cache — reverse lookup from blob key to format + // ═══════════════════════════════════════════════════════════════ + + /** + * Registers a blob key → format mapping in the reverse lookup cache. + * Called after a successful upload so that subsequent download/delete operations + * can resolve the correct format container for a plain blob key. + * + * @param blobKey the remote blob key (e.g., "_0.pqt__UUID") + * @param format the data format name (e.g., "parquet") + */ + public void registerBlobFormat(String blobKey, String format) { + if (blobKey != null && format != null) { + var updated = new HashMap<>(blobFormatCache); + updated.put(blobKey, format); + blobFormatCache = Map.copyOf(updated); + } + } + + /** + * Removes a blob key from the reverse lookup cache. + * Called after a file is deleted from the remote store. + * + * @param blobKey the remote blob key to unregister + */ + public void unregisterBlobFormat(String blobKey) { + if (blobKey != null) { + var updated = new HashMap<>(blobFormatCache); + updated.remove(blobKey); + blobFormatCache = Map.copyOf(updated); + } + } + + /** + * Atomically replaces the entire blob format cache. + * Called during initialization or full metadata refresh when the complete + * mapping is rebuilt from remote segment metadata. + * + * @param blobKeyToFormat the new complete mapping of blob key to format + */ + public void replaceBlobFormatCache(Map blobKeyToFormat) { + blobFormatCache = Map.copyOf(blobKeyToFormat); + } + + /** + * Resolves the data format for a plain blob key using the reverse lookup cache. + * Returns the default format ("lucene") if the key is not found. + * + * @param blobKey the blob key to resolve + * @return the data format name, defaults to "lucene" + */ + public String resolveFormat(String blobKey) { + String cached = blobFormatCache.get(blobKey); + return cached != null ? cached : DEFAULT_FORMAT; + } + + /** + * Clears the blob format cache. Called during close/cleanup. + */ + public void clearBlobFormatCache() { + blobFormatCache = Map.of(); + } + + private BlobContainer createFormatContainer(String format) { + BlobPath formatPath = basePath.add(format.toLowerCase(Locale.ROOT)); + return blobStore.blobContainer(formatPath); + } +} diff --git a/server/src/main/java/org/opensearch/indices/IndicesService.java b/server/src/main/java/org/opensearch/indices/IndicesService.java index 28593ca80dc83..7027574a0938f 100644 --- a/server/src/main/java/org/opensearch/indices/IndicesService.java +++ b/server/src/main/java/org/opensearch/indices/IndicesService.java @@ -426,6 +426,7 @@ public class IndicesService extends AbstractLifecycleComponent private final StatusCounterStats statusCounterStats; private final ClusterMergeSchedulerConfig clusterMergeSchedulerConfig; private final DataFormatRegistry dataFormatRegistry; + private final Map dataFormatAwareStoreDirectoryFactories; @Override protected void doStart() { @@ -454,6 +455,7 @@ public IndicesService( Collection>> engineFactoryProviders, Map directoryFactories, Map compositeDirectoryFactories, + Map dataFormatAwareStoreDirectoryFactories, ValuesSourceRegistry valuesSourceRegistry, Map recoveryStateFactories, Map storeFactories, @@ -469,7 +471,8 @@ public IndicesService( FileCache fileCache, CompositeIndexSettings compositeIndexSettings, Consumer replicator, - Function segmentReplicationStatsProvider + Function segmentReplicationStatsProvider, + DataFormatRegistry dataFormatRegistry ) { this.settings = settings; this.threadPool = threadPool; @@ -521,6 +524,7 @@ public void onRemoval(ShardId shardId, String fieldName, boolean wasEvicted, lon this.directoryFactories = directoryFactories; this.compositeDirectoryFactories = compositeDirectoryFactories; + this.dataFormatAwareStoreDirectoryFactories = dataFormatAwareStoreDirectoryFactories; this.recoveryStateFactories = recoveryStateFactories; this.storeFactories = storeFactories; this.ingestionConsumerFactories = ingestionConsumerFactories; @@ -611,7 +615,7 @@ protected void closeInternal() { MergeSchedulerConfig.CLUSTER_MAX_FORCE_MERGE_MB_PER_SEC_SETTING, this::onClusterLevelForceMergeMBPerSecUpdate ); - this.dataFormatRegistry = new DataFormatRegistry(pluginsService); + this.dataFormatRegistry = dataFormatRegistry; } @InternalApi @@ -665,6 +669,7 @@ public IndicesService( engineFactoryProviders, directoryFactories, Collections.emptyMap(), + Collections.emptyMap(), valuesSourceRegistry, recoveryStateFactories, Collections.emptyMap(), @@ -680,6 +685,7 @@ public IndicesService( null, null, null, + null, null ); } @@ -1103,7 +1109,8 @@ private synchronized IndexService createIndexService( recoveryStateFactories, storeFactories, fileCache, - compositeIndexSettings + compositeIndexSettings, + dataFormatAwareStoreDirectoryFactories ); for (IndexingOperationListener operationListener : indexingOperationListeners) { indexModule.addIndexOperationListener(operationListener); @@ -1226,7 +1233,8 @@ public synchronized MapperService createIndexMapperService(IndexMetadata indexMe recoveryStateFactories, storeFactories, fileCache, - compositeIndexSettings + compositeIndexSettings, + dataFormatAwareStoreDirectoryFactories ); pluginsService.onIndexModule(indexModule); return indexModule.newIndexMapperService(xContentRegistry, mapperRegistry, scriptService); diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java index f051abfffacf2..bde88570955da 100644 --- a/server/src/main/java/org/opensearch/node/Node.java +++ b/server/src/main/java/org/opensearch/node/Node.java @@ -164,6 +164,7 @@ import org.opensearch.index.compositeindex.CompositeIndexSettings; import org.opensearch.index.engine.EngineFactory; import org.opensearch.index.engine.MergedSegmentWarmerFactory; +import org.opensearch.index.engine.dataformat.DataFormatRegistry; import org.opensearch.index.mapper.MappingTransformerRegistry; import org.opensearch.index.recovery.RemoteStoreRestoreService; import org.opensearch.index.remote.RemoteIndexPathUploader; @@ -931,6 +932,14 @@ protected Node(final Environment initialEnvironment, Collection clas compositeDirectoryFactories.put(k, v); }); compositeDirectoryFactories.put("default", new DefaultCompositeDirectoryFactory()); + final Map dataFormatAwareStoreDirectoryFactories = + new HashMap<>(); + + // Register default factory + dataFormatAwareStoreDirectoryFactories.put( + "default", + new org.opensearch.index.store.DefaultDataFormatAwareStoreDirectoryFactory() + ); final Map recoveryStateFactories = pluginsService.filterPlugins( IndexStorePlugin.class @@ -955,10 +964,13 @@ protected Node(final Environment initialEnvironment, Collection clas final CompositeIndexSettings compositeIndexSettings = new CompositeIndexSettings(settings, settingsModule.getClusterSettings()); + final DataFormatRegistry dataFormatRegistry = new DataFormatRegistry(pluginsService); + final IndexStorePlugin.DirectoryFactory remoteDirectoryFactory = new RemoteSegmentStoreDirectoryFactory( repositoriesServiceReference::get, threadPool, - remoteStoreSettings.getSegmentsPathFixedPrefix() + remoteStoreSettings.getSegmentsPathFixedPrefix(), + dataFormatRegistry ); final TaskResourceTrackingService taskResourceTrackingService = new TaskResourceTrackingService( @@ -997,6 +1009,7 @@ protected Node(final Environment initialEnvironment, Collection clas engineFactoryProviders, Map.copyOf(directoryFactories), Map.copyOf(compositeDirectoryFactories), + Map.copyOf(dataFormatAwareStoreDirectoryFactories), searchModule.getValuesSourceRegistry(), recoveryStateFactories, storeFactories, @@ -1012,7 +1025,8 @@ protected Node(final Environment initialEnvironment, Collection clas fileCache, compositeIndexSettings, segmentReplicator::startReplication, - segmentReplicator::getSegmentReplicationStats + segmentReplicator::getSegmentReplicationStats, + dataFormatRegistry ); final IngestService ingestService = new IngestService( diff --git a/server/src/test/java/org/opensearch/index/IndexModuleTests.java b/server/src/test/java/org/opensearch/index/IndexModuleTests.java index ed1c61b4a5d15..c229d5c79e017 100644 --- a/server/src/test/java/org/opensearch/index/IndexModuleTests.java +++ b/server/src/test/java/org/opensearch/index/IndexModuleTests.java @@ -686,7 +686,8 @@ public void testStoreFactory() throws IOException { Collections.emptyMap(), storeFactories, null, - null + null, + Collections.emptyMap() ); // Test that IndexService can be created successfully with valid store factory @@ -718,7 +719,8 @@ public void testStoreFactoryWithEmptySetting() throws IOException { Collections.emptyMap(), storeFactories, null, - null + null, + Collections.emptyMap() ); // Test that IndexService uses default store when setting is empty @@ -748,7 +750,8 @@ public void testUnknownStoreFactory() { Collections.emptyMap(), Collections.emptyMap(), null, - null + null, + Collections.emptyMap() ); IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, () -> newIndexService(module)); diff --git a/server/src/test/java/org/opensearch/index/engine/dataformat/DataFormatPluginTests.java b/server/src/test/java/org/opensearch/index/engine/dataformat/DataFormatPluginTests.java index 87a34aaef255a..4f58e41f75c39 100644 --- a/server/src/test/java/org/opensearch/index/engine/dataformat/DataFormatPluginTests.java +++ b/server/src/test/java/org/opensearch/index/engine/dataformat/DataFormatPluginTests.java @@ -78,7 +78,8 @@ public void testFullDataFormatLifecycle() throws IOException { new ShardPath(false, Path.of("/tmp/uuid/0"), Path.of("/tmp/uuid/0"), new ShardId("index", "uuid", 0)), new IndexSettings(IndexMetadata.builder("index").settings(settings).build(), settings), null - ) + ), + null ); assertEquals(format, engine.getDataFormat()); diff --git a/server/src/test/java/org/opensearch/index/engine/dataformat/stub/MockCatalogSnapshot.java b/server/src/test/java/org/opensearch/index/engine/dataformat/stub/MockCatalogSnapshot.java index 1fa2a96a45b31..6f2ab935d3f41 100644 --- a/server/src/test/java/org/opensearch/index/engine/dataformat/stub/MockCatalogSnapshot.java +++ b/server/src/test/java/org/opensearch/index/engine/dataformat/stub/MockCatalogSnapshot.java @@ -74,7 +74,7 @@ public String serializeToString() { } @Override - public void setUserData(Map userData) {} + public void setUserData(Map userData, boolean commitData) {} @Override public CatalogSnapshot clone() { @@ -86,6 +86,21 @@ public void writeTo(StreamOutput out) throws IOException { super.writeTo(out); } + @Override + public int getFormatVersionForFile(String file) { + return 0; + } + + @Override + public byte[] serialize() throws IOException { + return new byte[0]; + } + + @Override + public Collection getFiles(boolean includeSegmentsFile) { + return List.of(); + } + @Override protected void closeInternal() {} } diff --git a/server/src/test/java/org/opensearch/index/engine/dataformat/stub/MockDataFormatPlugin.java b/server/src/test/java/org/opensearch/index/engine/dataformat/stub/MockDataFormatPlugin.java index 7a49fc6826ff9..da08e0e401b23 100644 --- a/server/src/test/java/org/opensearch/index/engine/dataformat/stub/MockDataFormatPlugin.java +++ b/server/src/test/java/org/opensearch/index/engine/dataformat/stub/MockDataFormatPlugin.java @@ -12,6 +12,7 @@ import org.opensearch.index.engine.dataformat.DataFormatPlugin; import org.opensearch.index.engine.dataformat.IndexingEngineConfig; import org.opensearch.index.engine.dataformat.IndexingExecutionEngine; +import org.opensearch.index.store.FormatChecksumStrategy; /** * A mock {@link DataFormatPlugin} for testing purposes. @@ -33,7 +34,7 @@ public DataFormat getDataFormat() { } @Override - public IndexingExecutionEngine indexingEngine(IndexingEngineConfig settings) { + public IndexingExecutionEngine indexingEngine(IndexingEngineConfig settings, FormatChecksumStrategy checksumStrategy) { return new MockIndexingExecutionEngine(dataFormat); } } 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 3b6089ab4729e..e43b5d3ece409 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 @@ -15,6 +15,7 @@ import java.io.IOException; import java.util.ArrayList; +import java.util.Collection; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; @@ -345,6 +346,70 @@ public void testTryIncRefFailsAfterClosed() { assertFalse(snapshot.tryIncRef()); } + public void testGetUploadFileNamesProducesFormatSlashFile() throws Exception { + // Build a snapshot with known segments and files + WriterFileSet parquetWfs = new WriterFileSet("/tmp/pq", 1L, Set.of("_0.pqt", "_1.pqt"), 100); + WriterFileSet luceneWfs = new WriterFileSet("/tmp/lc", 1L, Set.of("_0.cfe", "_0.si"), 50); + Segment segment = new Segment(1L, Map.of("parquet", parquetWfs, "lucene", luceneWfs)); + DataformatAwareCatalogSnapshot snapshot = new DataformatAwareCatalogSnapshot(1L, 1L, 1L, List.of(segment), 0L, Map.of()); + + Collection uploadNames = snapshot.getFiles(true); + + // Parquet files: "parquet/_0.pqt", "parquet/_1.pqt" + // Lucene files: plain names "_0.cfe", "_0.si" (FileMetadata.serialize() omits "lucene/" prefix) + assertEquals(4, uploadNames.size()); + assertTrue(uploadNames.contains("parquet/_0.pqt")); + assertTrue(uploadNames.contains("parquet/_1.pqt")); + assertTrue(uploadNames.contains("_0.cfe")); + assertTrue(uploadNames.contains("_0.si")); + } + + public void testGetFilesEmptySegments() throws Exception { + DataformatAwareCatalogSnapshot snapshot = new DataformatAwareCatalogSnapshot(1L, 1L, 1L, List.of(), 0L, Map.of()); + Collection uploadNames = snapshot.getFiles(true); + assertTrue(uploadNames.isEmpty()); + } + + public void testGetFilesMultipleSegments() throws Exception { + WriterFileSet wfs1 = new WriterFileSet("/tmp/pq", 1L, Set.of("_0.pqt"), 10); + WriterFileSet wfs2 = new WriterFileSet("/tmp/pq", 2L, Set.of("_1.pqt"), 20); + Segment seg1 = new Segment(1L, Map.of("parquet", wfs1)); + Segment seg2 = new Segment(2L, Map.of("parquet", wfs2)); + DataformatAwareCatalogSnapshot snapshot = new DataformatAwareCatalogSnapshot(1L, 1L, 1L, List.of(seg1, seg2), 0L, Map.of()); + + Collection uploadNames = snapshot.getFiles(true); + assertEquals(2, uploadNames.size()); + assertTrue(uploadNames.contains("parquet/_0.pqt")); + assertTrue(uploadNames.contains("parquet/_1.pqt")); + } + + public void testSerializeThrowsUnsupportedOperation() { + DataformatAwareCatalogSnapshot snapshot = new DataformatAwareCatalogSnapshot(1L, 1L, 1L, List.of(), 0L, Map.of()); + expectThrows(UnsupportedOperationException.class, snapshot::serialize); + } + + public void testGetFormatVersionForFileReturnsOpenSearchMajor() { + DataformatAwareCatalogSnapshot snapshot = new DataformatAwareCatalogSnapshot(1L, 1L, 1L, List.of(), 0L, Map.of()); + assertEquals(org.opensearch.Version.CURRENT.major, snapshot.getFormatVersionForFile("any_file.pqt")); + } + + public void testSetUserDataUpdatesAndReturns() { + DataformatAwareCatalogSnapshot snapshot = new DataformatAwareCatalogSnapshot(1L, 1L, 1L, List.of(), 0L, Map.of("a", "b")); + assertEquals(Map.of("a", "b"), snapshot.getUserData()); + snapshot.setUserData(Map.of("x", "y"), false); + assertEquals(Map.of("x", "y"), snapshot.getUserData()); + } + + public void testClonePreservesUserData() { + Map userData = Map.of("key1", "val1", "key2", "val2"); + DataformatAwareCatalogSnapshot snapshot = new DataformatAwareCatalogSnapshot(42L, 10L, 5L, List.of(), 3L, userData); + DataformatAwareCatalogSnapshot cloned = snapshot.clone(); + assertEquals(userData, cloned.getUserData()); + assertEquals(42L, cloned.getId()); + assertEquals(10L, cloned.getGeneration()); + assertEquals(5L, cloned.getVersion()); + } + // --- helpers --- private WriterFileSet randomWriterFileSet(String format) { diff --git a/server/src/test/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshotTests.java b/server/src/test/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshotTests.java index 0c8986d3a1043..b627d07cb151c 100644 --- a/server/src/test/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshotTests.java +++ b/server/src/test/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshotTests.java @@ -13,6 +13,7 @@ import org.opensearch.core.common.io.stream.NamedWriteableRegistry; import org.opensearch.test.OpenSearchTestCase; +import java.util.Collection; import java.util.Collections; import java.util.HashMap; import java.util.Map; @@ -40,7 +41,7 @@ public void testDelegation() { expectThrows(UnsupportedOperationException.class, snapshot::serializeToString); // setUserData is a no-op, should not throw - snapshot.setUserData(Map.of("key", "value")); + snapshot.setUserData(Map.of("key", "value"), false); } } @@ -51,7 +52,7 @@ public void testClone() { SegmentInfosCatalogSnapshot cloned = snapshot.clone(); assertNotSame(snapshot, cloned); - assertSame(segmentInfos, cloned.getSegmentInfos()); + assertNotSame(segmentInfos, cloned.getSegmentInfos()); assertEquals(snapshot.getGeneration(), cloned.getGeneration()); assertEquals(snapshot.getVersion(), cloned.getVersion()); } @@ -72,6 +73,48 @@ public void testCopyWriteable() throws Exception { assertEquals(original.getUserData(), copy.getUserData()); } + public void testGetFiles() throws Exception { + SegmentInfos segmentInfos = randomSegmentInfos(); + SegmentInfosCatalogSnapshot snapshot = new SegmentInfosCatalogSnapshot(segmentInfos); + Collection uploadNames = snapshot.getFiles(true); + assertEquals(segmentInfos.files(true), new java.util.HashSet<>(uploadNames)); + } + + public void testSerializeProducesValidBytes() throws Exception { + SegmentInfos segmentInfos = randomSegmentInfos(); + SegmentInfosCatalogSnapshot snapshot = new SegmentInfosCatalogSnapshot(segmentInfos); + byte[] bytes = snapshot.serialize(); + assertNotNull(bytes); + assertTrue(bytes.length > 0); + } + + public void testGetFormatVersionForUnmappedFileDefaultsToLatest() { + SegmentInfos segmentInfos = randomSegmentInfos(); + SegmentInfosCatalogSnapshot snapshot = new SegmentInfosCatalogSnapshot(segmentInfos); + // File not in segmentFileVersionMap and not the segments file → falls through to LATEST.major + int version = snapshot.getFormatVersionForFile("nonexistent_file.xyz"); + assertEquals(Version.LATEST.major, version); + } + + public void testSetUserDataDelegatesToSegmentInfos() { + SegmentInfos segmentInfos = randomSegmentInfos(); + SegmentInfosCatalogSnapshot snapshot = new SegmentInfosCatalogSnapshot(segmentInfos); + Map newData = Map.of("key1", "val1", "key2", "val2"); + snapshot.setUserData(newData, false); + assertEquals(newData, segmentInfos.getUserData()); + assertEquals(newData, snapshot.getUserData()); + } + + public void testCloneNoAcquireReturnsIndependentCopy() { + SegmentInfos segmentInfos = randomSegmentInfos(); + SegmentInfosCatalogSnapshot snapshot = new SegmentInfosCatalogSnapshot(segmentInfos); + CatalogSnapshot cloned = snapshot.cloneNoAcquire(); + assertNotSame(snapshot, cloned); + assertNotSame(segmentInfos, ((SegmentInfosCatalogSnapshot) cloned).getSegmentInfos()); + assertEquals(snapshot.getGeneration(), cloned.getGeneration()); + assertEquals(snapshot.getVersion(), cloned.getVersion()); + } + // --- helpers --- private SegmentInfos randomSegmentInfos() { diff --git a/server/src/test/java/org/opensearch/index/shard/RemoteStoreRefreshListenerTests.java b/server/src/test/java/org/opensearch/index/shard/RemoteStoreRefreshListenerTests.java index cb451ab6761a9..02a8d02c3b558 100644 --- a/server/src/test/java/org/opensearch/index/shard/RemoteStoreRefreshListenerTests.java +++ b/server/src/test/java/org/opensearch/index/shard/RemoteStoreRefreshListenerTests.java @@ -29,11 +29,11 @@ import org.opensearch.core.index.shard.ShardId; import org.opensearch.index.engine.InternalEngineFactory; import org.opensearch.index.engine.NRTReplicationEngineFactory; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import org.opensearch.index.remote.RemoteSegmentTransferTracker; import org.opensearch.index.remote.RemoteStoreStatsTrackerFactory; import org.opensearch.index.store.RemoteDirectory; import org.opensearch.index.store.RemoteSegmentStoreDirectory; -import org.opensearch.index.store.RemoteSegmentStoreDirectory.MetadataFilenameUtils; import org.opensearch.index.store.Store; import org.opensearch.index.store.lockmanager.RemoteStoreLockManager; import org.opensearch.indices.DefaultRemoteStoreSettings; @@ -213,7 +213,10 @@ public void testRemoteDirectoryInitThrowsException() throws IOException { } throw new IOException(); }).when(remoteMetadataDirectory) - .listFilesByPrefixInLexicographicOrder(MetadataFilenameUtils.METADATA_PREFIX, METADATA_FILES_TO_FETCH); + .listFilesByPrefixInLexicographicOrder( + RemoteSegmentStoreDirectory.MetadataFilenameUtils.METADATA_PREFIX, + METADATA_FILES_TO_FETCH + ); SegmentInfos segmentInfos; try (Store indexShardStore = indexShard.store()) { @@ -248,7 +251,7 @@ public void testRemoteDirectoryInitThrowsException() throws IOException { // listFilesByPrefixInLexicographicOrder has been called twice. verify(remoteMetadataDirectory, times(1)).getBlobStream(any()); verify(remoteMetadataDirectory, times(2)).listFilesByPrefixInLexicographicOrder( - MetadataFilenameUtils.METADATA_PREFIX, + RemoteSegmentStoreDirectory.MetadataFilenameUtils.METADATA_PREFIX, METADATA_FILES_TO_FETCH ); } @@ -767,21 +770,21 @@ private Tuple mockIn }).when(shard).isStartedPrimary(); AtomicLong counter = new AtomicLong(); - // Mock indexShard.getSegmentInfosSnapshot() + // Mock indexShard.getCatalogSnapshot() doAnswer(invocation -> { if (counter.incrementAndGet() <= succeedOnAttempt) { - logger.error("Failing in get segment info {}", counter.get()); + logger.error("Failing in get catalog snapshot {}", counter.get()); throw new RuntimeException("Inducing failure in upload"); } - return indexShard.getSegmentInfosSnapshot(); - }).when(shard).getSegmentInfosSnapshot(); + return indexShard.getCatalogSnapshot(); + }).when(shard).getCatalogSnapshot(); doAnswer((invocation -> { if (counter.incrementAndGet() <= succeedOnAttempt) { throw new RuntimeException("Inducing failure in upload"); } return indexShard.getLatestReplicationCheckpoint(); - })).when(shard).computeReplicationCheckpoint(any()); + })).when(shard).computeReplicationCheckpoint(any(CatalogSnapshot.class)); doAnswer((invocationOnMock -> { if (closeShard && counter.get() == closeShardAfterAttempt) { diff --git a/server/src/test/java/org/opensearch/index/shard/SegmentReplicationIndexShardTests.java b/server/src/test/java/org/opensearch/index/shard/SegmentReplicationIndexShardTests.java index 45efaec109f90..fb9f6e1095748 100644 --- a/server/src/test/java/org/opensearch/index/shard/SegmentReplicationIndexShardTests.java +++ b/server/src/test/java/org/opensearch/index/shard/SegmentReplicationIndexShardTests.java @@ -1152,7 +1152,7 @@ public void testReuseReplicationCheckpointWhenLatestInfosIsUnChanged() throws Ex public void testComputeReplicationCheckpointNullInfosReturnsEmptyCheckpoint() throws Exception { try (ReplicationGroup shards = createGroup(1, settings, indexMapping, new NRTReplicationEngineFactory(), createTempDir())) { final IndexShard primaryShard = shards.getPrimary(); - assertEquals(ReplicationCheckpoint.empty(primaryShard.shardId), primaryShard.computeReplicationCheckpoint(null)); + assertEquals(ReplicationCheckpoint.empty(primaryShard.shardId), primaryShard.computeReplicationCheckpoint((SegmentInfos) null)); } } diff --git a/server/src/test/java/org/opensearch/index/store/BaseRemoteSegmentStoreDirectoryTests.java b/server/src/test/java/org/opensearch/index/store/BaseRemoteSegmentStoreDirectoryTests.java index 1c5dec20fd74a..0d7cb72c137c2 100644 --- a/server/src/test/java/org/opensearch/index/store/BaseRemoteSegmentStoreDirectoryTests.java +++ b/server/src/test/java/org/opensearch/index/store/BaseRemoteSegmentStoreDirectoryTests.java @@ -44,7 +44,15 @@ public class BaseRemoteSegmentStoreDirectoryTests extends IndexShardTestCase { protected SegmentInfos segmentInfos; protected ThreadPool threadPool; - protected String metadataFilename = RemoteSegmentStoreDirectory.MetadataFilenameUtils.getMetadataFilename(12, 23, 34, 1, 1, "node-1"); + protected String metadataFilename = RemoteSegmentStoreDirectory.MetadataFilenameUtils.getMetadataFilename( + 12, + 23, + 34, + 1, + 1, + "node-1", + 0L + ); protected String metadataFilenameDup = RemoteSegmentStoreDirectory.MetadataFilenameUtils.getMetadataFilename( 12, @@ -52,11 +60,36 @@ public class BaseRemoteSegmentStoreDirectoryTests extends IndexShardTestCase { 34, 2, 1, - "node-2" + "node-2", + 0L + ); + protected String metadataFilename2 = RemoteSegmentStoreDirectory.MetadataFilenameUtils.getMetadataFilename( + 12, + 13, + 34, + 1, + 1, + "node-1", + 0L + ); + protected String metadataFilename3 = RemoteSegmentStoreDirectory.MetadataFilenameUtils.getMetadataFilename( + 10, + 38, + 34, + 1, + 1, + "node-1", + 0L + ); + protected String metadataFilename4 = RemoteSegmentStoreDirectory.MetadataFilenameUtils.getMetadataFilename( + 10, + 36, + 34, + 1, + 1, + "node-1", + 0L ); - protected String metadataFilename2 = RemoteSegmentStoreDirectory.MetadataFilenameUtils.getMetadataFilename(12, 13, 34, 1, 1, "node-1"); - protected String metadataFilename3 = RemoteSegmentStoreDirectory.MetadataFilenameUtils.getMetadataFilename(10, 38, 34, 1, 1, "node-1"); - protected String metadataFilename4 = RemoteSegmentStoreDirectory.MetadataFilenameUtils.getMetadataFilename(10, 36, 34, 1, 1, "node-1"); public void setupRemoteSegmentStoreDirectory() throws IOException { remoteDataDirectory = mock(RemoteDirectory.class); diff --git a/server/src/test/java/org/opensearch/index/store/DataFormatAwareStoreDirectoryTests.java b/server/src/test/java/org/opensearch/index/store/DataFormatAwareStoreDirectoryTests.java new file mode 100644 index 0000000000000..ba795396451b5 --- /dev/null +++ b/server/src/test/java/org/opensearch/index/store/DataFormatAwareStoreDirectoryTests.java @@ -0,0 +1,992 @@ +/* + * 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.store; + +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.store.FSDirectory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.IndexOutput; +import org.opensearch.Version; +import org.opensearch.cluster.metadata.IndexMetadata; +import org.opensearch.common.settings.Settings; +import org.opensearch.core.index.Index; +import org.opensearch.core.index.shard.ShardId; +import org.opensearch.index.IndexSettings; +import org.opensearch.index.engine.dataformat.DataFormatPlugin; +import org.opensearch.index.engine.dataformat.DataFormatRegistry; +import org.opensearch.index.shard.ShardPath; +import org.opensearch.plugins.PluginsService; +import org.opensearch.plugins.SearchBackEndPlugin; +import org.opensearch.test.OpenSearchTestCase; +import org.junit.After; +import org.junit.Before; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.List; +import java.util.Set; +import java.util.zip.CRC32; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class DataFormatAwareStoreDirectoryTests extends OpenSearchTestCase { + + private Path tempDir; + private Path shardDataPath; + private Path indexPath; + private FSDirectory fsDirectory; + private ShardPath shardPath; + private DataFormatAwareStoreDirectory dataFormatAwareStoreDirectory; + + @Before + public void setUp() throws Exception { + super.setUp(); + // Create directory structure: tempDir///index/ + tempDir = createTempDir(); + String indexUUID = "test-index-uuid"; + int shardId = 0; + shardDataPath = tempDir.resolve(indexUUID).resolve(Integer.toString(shardId)); + indexPath = shardDataPath.resolve(ShardPath.INDEX_FOLDER_NAME); + Files.createDirectories(indexPath); + + fsDirectory = FSDirectory.open(indexPath); + ShardId sid = new ShardId(new Index("test-index", indexUUID), shardId); + shardPath = new ShardPath(false, shardDataPath, shardDataPath, sid); + + PluginsService pluginsService = mock(PluginsService.class); + when(pluginsService.filterPlugins(DataFormatPlugin.class)).thenReturn(List.of()); + when(pluginsService.filterPlugins(SearchBackEndPlugin.class)).thenReturn(List.of()); + DataFormatRegistry dataFormatRegistry = new DataFormatRegistry(pluginsService); + + // Create real IndexSettings (IndexSettings is final, cannot be mocked) + Settings settings = Settings.builder() + .put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT) + .put(IndexMetadata.SETTING_INDEX_UUID, indexUUID) + .build(); + IndexMetadata metadata = IndexMetadata.builder("test-index").settings(settings).numberOfShards(1).numberOfReplicas(0).build(); + IndexSettings indexSettings = new IndexSettings(metadata, Settings.EMPTY); + + dataFormatAwareStoreDirectory = new DataFormatAwareStoreDirectory(indexSettings, fsDirectory, shardPath, dataFormatRegistry); + } + + @After + public void tearDown() throws Exception { + if (dataFormatAwareStoreDirectory != null) { + dataFormatAwareStoreDirectory.close(); + } + super.tearDown(); + } + + // ═══════════════════════════════════════════════════════════════ + // toFileMetadata / toFileIdentifier + // ═══════════════════════════════════════════════════════════════ + + public void testToFileMetadata_luceneFile() { + FileMetadata fm = DataFormatAwareStoreDirectory.toFileMetadata("_0.si"); + assertEquals("lucene", fm.dataFormat()); + assertEquals("_0.si", fm.file()); + } + + public void testToFileMetadata_prefixedFile() { + FileMetadata fm = DataFormatAwareStoreDirectory.toFileMetadata("parquet/data.parquet"); + assertEquals("parquet", fm.dataFormat()); + assertEquals("data.parquet", fm.file()); + } + + public void testToFileMetadata_arrowFile() { + FileMetadata fm = DataFormatAwareStoreDirectory.toFileMetadata("arrow/data.arrow"); + assertEquals("arrow", fm.dataFormat()); + assertEquals("data.arrow", fm.file()); + } + + public void testToFileIdentifier_lucene() { + FileMetadata fm = new FileMetadata("lucene", "_0.si"); + String identifier = DataFormatAwareStoreDirectory.toFileIdentifier(fm); + assertEquals("_0.si", identifier); + } + + public void testToFileIdentifier_metadata() { + // "metadata" is treated as a default/index format, so no prefix is added + FileMetadata fm = new FileMetadata("metadata", "meta_file.txt"); + String identifier = DataFormatAwareStoreDirectory.toFileIdentifier(fm); + assertEquals("meta_file.txt", identifier); + } + + public void testToFileIdentifier_nonLucene() { + FileMetadata fm = new FileMetadata("parquet", "data.parquet"); + String identifier = DataFormatAwareStoreDirectory.toFileIdentifier(fm); + assertEquals("parquet/data.parquet", identifier); + } + + public void testToFileIdentifier_arrow() { + FileMetadata fm = new FileMetadata("arrow", "data.arrow"); + String identifier = DataFormatAwareStoreDirectory.toFileIdentifier(fm); + assertEquals("arrow/data.arrow", identifier); + } + + public void testRoundtrip_toFileMetadata_toFileIdentifier_lucene() { + String original = "_0.cfe"; + FileMetadata fm = DataFormatAwareStoreDirectory.toFileMetadata(original); + String result = DataFormatAwareStoreDirectory.toFileIdentifier(fm); + assertEquals(original, result); + } + + public void testRoundtrip_toFileMetadata_toFileIdentifier_nonLucene() { + String original = "parquet/data.parquet"; + FileMetadata fm = DataFormatAwareStoreDirectory.toFileMetadata(original); + String result = DataFormatAwareStoreDirectory.toFileIdentifier(fm); + assertEquals(original, result); + } + + // ═══════════════════════════════════════════════════════════════ + // getDataFormat + // ═══════════════════════════════════════════════════════════════ + + public void testGetDataFormat_lucene() { + assertEquals("lucene", dataFormatAwareStoreDirectory.getDataFormat("_0.cfe")); + } + + public void testGetDataFormat_nonLucene() { + assertEquals("arrow", dataFormatAwareStoreDirectory.getDataFormat("arrow/data.arrow")); + } + + public void testGetDataFormat_parquet() { + assertEquals("parquet", dataFormatAwareStoreDirectory.getDataFormat("parquet/data.parquet")); + } + + // ═══════════════════════════════════════════════════════════════ + // createOutput / openInput - Lucene files + // ═══════════════════════════════════════════════════════════════ + + public void testCreateOutputAndOpenInput_lucene() throws IOException { + String fileName = "_0_test.si"; + byte[] testData = "hello world lucene".getBytes(StandardCharsets.UTF_8); + + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fileName, IOContext.DEFAULT)) { + out.writeBytes(testData, testData.length); + } + + try (IndexInput in = dataFormatAwareStoreDirectory.openInput(fileName, IOContext.DEFAULT)) { + byte[] readData = new byte[testData.length]; + in.readBytes(readData, 0, readData.length); + assertArrayEquals(testData, readData); + } + } + + // ═══════════════════════════════════════════════════════════════ + // createOutput / openInput - Non-Lucene files + // ═══════════════════════════════════════════════════════════════ + + public void testCreateOutputAndOpenInput_nonLucene() throws IOException { + String fileIdentifier = "parquet/data.parquet"; + byte[] testData = "hello world parquet".getBytes(StandardCharsets.UTF_8); + + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fileIdentifier, IOContext.DEFAULT)) { + out.writeBytes(testData, testData.length); + } + + try (IndexInput in = dataFormatAwareStoreDirectory.openInput(fileIdentifier, IOContext.DEFAULT)) { + byte[] readData = new byte[testData.length]; + in.readBytes(readData, 0, readData.length); + assertArrayEquals(testData, readData); + } + } + + // ═══════════════════════════════════════════════════════════════ + // fileLength + // ═══════════════════════════════════════════════════════════════ + + public void testFileLength_lucene() throws IOException { + String fileName = "_test_len.si"; + byte[] data = "some content for length test".getBytes(StandardCharsets.UTF_8); + + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fileName, IOContext.DEFAULT)) { + out.writeBytes(data, data.length); + } + + assertEquals(data.length, dataFormatAwareStoreDirectory.fileLength(fileName)); + } + + public void testFileLength_fileMetadata() throws IOException { + String fileIdentifier = "parquet/len_test.parquet"; + byte[] data = "parquet length test".getBytes(StandardCharsets.UTF_8); + + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fileIdentifier, IOContext.DEFAULT)) { + out.writeBytes(data, data.length); + } + + FileMetadata fm = new FileMetadata("parquet", "len_test.parquet"); + assertEquals(data.length, dataFormatAwareStoreDirectory.fileLength(fm.serialize())); + } + + // ═══════════════════════════════════════════════════════════════ + // deleteFile + // ═══════════════════════════════════════════════════════════════ + + public void testDeleteFile_string() throws IOException { + String fileName = "_del_test.si"; + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fileName, IOContext.DEFAULT)) { + out.writeString("to be deleted"); + } + assertTrue(Arrays.asList(dataFormatAwareStoreDirectory.listAll()).contains(fileName)); + + dataFormatAwareStoreDirectory.deleteFile(fileName); + assertFalse(Arrays.asList(dataFormatAwareStoreDirectory.listAll()).contains(fileName)); + } + + public void testDeleteFile_fileMetadata() throws IOException { + String fileIdentifier = "parquet/del_test.parquet"; + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fileIdentifier, IOContext.DEFAULT)) { + out.writeString("to be deleted"); + } + String serialized = new FileMetadata("parquet", "del_test.parquet").serialize(); + assertTrue(Arrays.asList(dataFormatAwareStoreDirectory.listAll()).contains(serialized)); + + FileMetadata fm = new FileMetadata("parquet", "del_test.parquet"); + dataFormatAwareStoreDirectory.deleteFile(fm.serialize()); + assertFalse(Arrays.asList(dataFormatAwareStoreDirectory.listAll()).contains(serialized)); + } + + // ═══════════════════════════════════════════════════════════════ + // listAll / listFileMetadata + // ═══════════════════════════════════════════════════════════════ + + public void testListAll_empty() throws IOException { + // A fresh directory should list no segment files + String[] files = dataFormatAwareStoreDirectory.listAll(); + assertNotNull(files); + } + + public void testListAll_withFiles() throws IOException { + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput("_0.si", IOContext.DEFAULT)) { + out.writeString("data"); + } + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput("parquet/data.parquet", IOContext.DEFAULT)) { + out.writeString("parquet data"); + } + + String[] files = dataFormatAwareStoreDirectory.listAll(); + List fileList = Arrays.asList(files); + assertTrue(fileList.contains("_0.si")); + assertTrue(fileList.contains("parquet/data.parquet")); + } + + // ═══════════════════════════════════════════════════════════════ + // Checksum - Lucene file (CodecUtil path) + // ═══════════════════════════════════════════════════════════════ + + public void testCalculateChecksum_luceneFile() throws IOException { + String fileName = "_cksum.si"; + // Write a file with Lucene CodecUtil header/footer so CodecUtil.retrieveChecksum works + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fileName, IOContext.DEFAULT)) { + CodecUtil.writeHeader(out, "TestCodec", 1); + out.writeString("some segment data for checksum"); + CodecUtil.writeFooter(out); + } + + long checksum = dataFormatAwareStoreDirectory.calculateChecksum(fileName); + // Verify it's a valid non-zero checksum (CodecUtil stores checksum in footer) + assertTrue("Checksum should be a valid value", checksum != 0); + + // Verify we get the same checksum via string overload with serialized name + FileMetadata fm = DataFormatAwareStoreDirectory.toFileMetadata(fileName); + assertEquals(checksum, dataFormatAwareStoreDirectory.calculateChecksum(fm.serialize())); + } + + // ═══════════════════════════════════════════════════════════════ + // Checksum - Non-Lucene file (CRC32 path) + // ═══════════════════════════════════════════════════════════════ + + public void testCalculateChecksum_nonLuceneFile() throws IOException { + String fileIdentifier = "parquet/cksum.parquet"; + byte[] data = "parquet content for checksum".getBytes(StandardCharsets.UTF_8); + + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fileIdentifier, IOContext.DEFAULT)) { + out.writeBytes(data, data.length); + } + + long checksum = dataFormatAwareStoreDirectory.calculateChecksum(fileIdentifier); + + // Compute expected CRC32 manually + CRC32 crc32 = new CRC32(); + crc32.update(data); + assertEquals("CRC32 checksum should match", crc32.getValue(), checksum); + } + + // ═══════════════════════════════════════════════════════════════ + // calculateUploadChecksum + // ═══════════════════════════════════════════════════════════════ + + public void testCalculateUploadChecksum() throws IOException { + String fileIdentifier = "parquet/upload_cksum.parquet"; + byte[] data = "upload checksum test".getBytes(StandardCharsets.UTF_8); + + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fileIdentifier, IOContext.DEFAULT)) { + out.writeBytes(data, data.length); + } + + FileMetadata fm = new FileMetadata("parquet", "upload_cksum.parquet"); + String uploadChecksum = dataFormatAwareStoreDirectory.calculateUploadChecksum(fm.serialize()); + assertNotNull(uploadChecksum); + // Should be the string representation of the long checksum + long parsedChecksum = Long.parseLong(uploadChecksum); + assertEquals(dataFormatAwareStoreDirectory.calculateChecksum(fm.serialize()), parsedChecksum); + } + + // ═══════════════════════════════════════════════════════════════ + // rename + // ═══════════════════════════════════════════════════════════════ + + public void testRename_sameFormat() throws IOException { + String fileName = "_rename_src.si"; + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fileName, IOContext.DEFAULT)) { + out.writeString("rename test data"); + } + + dataFormatAwareStoreDirectory.rename(fileName, "_rename_dest.si"); + assertFalse(Arrays.asList(dataFormatAwareStoreDirectory.listAll()).contains(fileName)); + assertTrue(Arrays.asList(dataFormatAwareStoreDirectory.listAll()).contains("_rename_dest.si")); + } + + public void testRename_fileMetadata_sameFormat() throws IOException { + String fileIdentifier = "parquet/rename_src.parquet"; + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fileIdentifier, IOContext.DEFAULT)) { + out.writeString("rename test data parquet"); + } + + FileMetadata src = new FileMetadata("parquet", "rename_src.parquet"); + FileMetadata dest = new FileMetadata("parquet", "rename_dest.parquet"); + dataFormatAwareStoreDirectory.rename(src.serialize(), dest.serialize()); + + String srcSerialized = new FileMetadata("parquet", "rename_src.parquet").serialize(); + String destSerialized = new FileMetadata("parquet", "rename_dest.parquet").serialize(); + assertFalse(Arrays.asList(dataFormatAwareStoreDirectory.listAll()).contains(srcSerialized)); + assertTrue(Arrays.asList(dataFormatAwareStoreDirectory.listAll()).contains(destSerialized)); + } + + // ═══════════════════════════════════════════════════════════════ + // sync + // ═══════════════════════════════════════════════════════════════ + + public void testSync() throws IOException { + String fileName = "_sync_test.si"; + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fileName, IOContext.DEFAULT)) { + out.writeString("sync test data"); + } + + // sync should not throw + dataFormatAwareStoreDirectory.sync(Set.of(fileName)); + } + + public void testGetShardPath() { + assertNotNull(dataFormatAwareStoreDirectory.getShardPath()); + assertEquals(shardPath, dataFormatAwareStoreDirectory.getShardPath()); + } + + // ═══════════════════════════════════════════════════════════════ + // FileMetadata convenience: getChecksumOfLocalFile + // ═══════════════════════════════════════════════════════════════ + + public void testGetChecksumOfLocalFile() throws IOException { + String fileIdentifier = "parquet/local_cksum.parquet"; + byte[] data = "local checksum test".getBytes(StandardCharsets.UTF_8); + + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fileIdentifier, IOContext.DEFAULT)) { + out.writeBytes(data, data.length); + } + + FileMetadata fm = new FileMetadata("parquet", "local_cksum.parquet"); + long checksum = dataFormatAwareStoreDirectory.calculateChecksum(fm.serialize()); + assertTrue(checksum != 0); + } + + // ═══════════════════════════════════════════════════════════════ + // resolveFileName with FileMetadata.DELIMITER in name + // ═══════════════════════════════════════════════════════════════ + + public void testOpenInput_withSerializedFileMetadata() throws IOException { + // Write a file using plain identifier + String fileIdentifier = "parquet/delimited_test.parquet"; + byte[] data = "delimited test data".getBytes(StandardCharsets.UTF_8); + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fileIdentifier, IOContext.DEFAULT)) { + out.writeBytes(data, data.length); + } + + // Now access it using a serialized FileMetadata string (with "/" delimiter) + FileMetadata fm = new FileMetadata("parquet", "delimited_test.parquet"); + String serialized = fm.serialize(); // "parquet/delimited_test.parquet" + + // The resolveFileName method should handle the delimiter and resolve correctly + long length = dataFormatAwareStoreDirectory.fileLength(serialized); + assertEquals(data.length, length); + } + + // ═══════════════════════════════════════════════════════════════ + // FileMetadata-based openInput / createOutput + // ═══════════════════════════════════════════════════════════════ + + public void testCreateOutputAndOpenInput_fileMetadata_lucene() throws IOException { + FileMetadata fm = new FileMetadata("lucene", "_fm_test.si"); + byte[] testData = "file metadata lucene test".getBytes(StandardCharsets.UTF_8); + + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fm, IOContext.DEFAULT)) { + out.writeBytes(testData, testData.length); + } + + try (IndexInput in = dataFormatAwareStoreDirectory.openInput(fm.serialize(), IOContext.DEFAULT)) { + byte[] readData = new byte[testData.length]; + in.readBytes(readData, 0, readData.length); + assertArrayEquals(testData, readData); + } + } + + public void testCreateOutputAndOpenInput_fileMetadata_parquet() throws IOException { + FileMetadata fm = new FileMetadata("parquet", "fm_test.parquet"); + byte[] testData = "file metadata parquet test".getBytes(StandardCharsets.UTF_8); + + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fm, IOContext.DEFAULT)) { + out.writeBytes(testData, testData.length); + } + + try (IndexInput in = dataFormatAwareStoreDirectory.openInput(fm.serialize(), IOContext.DEFAULT)) { + byte[] readData = new byte[testData.length]; + in.readBytes(readData, 0, readData.length); + assertArrayEquals(testData, readData); + } + } + + // ═══════════════════════════════════════════════════════════════ + // Checksum idempotency + // ═══════════════════════════════════════════════════════════════ + + public void testCalculateChecksum_idempotent_nonLucene() throws IOException { + String fileIdentifier = "parquet/idempotent.parquet"; + byte[] data = "idempotent checksum data".getBytes(StandardCharsets.UTF_8); + + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fileIdentifier, IOContext.DEFAULT)) { + out.writeBytes(data, data.length); + } + + long checksum1 = dataFormatAwareStoreDirectory.calculateChecksum(fileIdentifier); + long checksum2 = dataFormatAwareStoreDirectory.calculateChecksum(fileIdentifier); + assertEquals("Checksum should be the same on repeated calls", checksum1, checksum2); + } + + public void testCalculateChecksum_stringAndFileMetadataConsistent() throws IOException { + String fileIdentifier = "parquet/consistency.parquet"; + byte[] data = "consistency checksum".getBytes(StandardCharsets.UTF_8); + + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fileIdentifier, IOContext.DEFAULT)) { + out.writeBytes(data, data.length); + } + + long checksumFromString = dataFormatAwareStoreDirectory.calculateChecksum(fileIdentifier); + assertEquals(checksumFromString, dataFormatAwareStoreDirectory.calculateChecksum("parquet/consistency.parquet")); + } + + // ═══════════════════════════════════════════════════════════════ + // Data integrity after rename + // ═══════════════════════════════════════════════════════════════ + + public void testRename_preservesContent_lucene() throws IOException { + String srcFile = "_rename_content.si"; + byte[] data = "content to preserve after rename".getBytes(StandardCharsets.UTF_8); + + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(srcFile, IOContext.DEFAULT)) { + out.writeBytes(data, data.length); + } + + dataFormatAwareStoreDirectory.rename(srcFile, "_rename_content_dest.si"); + + try (IndexInput in = dataFormatAwareStoreDirectory.openInput("_rename_content_dest.si", IOContext.DEFAULT)) { + byte[] readData = new byte[data.length]; + in.readBytes(readData, 0, readData.length); + assertArrayEquals(data, readData); + } + } + + public void testRename_preservesContent_nonLucene() throws IOException { + FileMetadata src = new FileMetadata("parquet", "rename_content.parquet"); + FileMetadata dest = new FileMetadata("parquet", "rename_content_dest.parquet"); + byte[] data = "parquet content to preserve".getBytes(StandardCharsets.UTF_8); + + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(src, IOContext.DEFAULT)) { + out.writeBytes(data, data.length); + } + + dataFormatAwareStoreDirectory.rename(src.serialize(), dest.serialize()); + + try (IndexInput in = dataFormatAwareStoreDirectory.openInput(dest.serialize(), IOContext.DEFAULT)) { + byte[] readData = new byte[data.length]; + in.readBytes(readData, 0, readData.length); + assertArrayEquals(data, readData); + } + } + + // ═══════════════════════════════════════════════════════════════ + // Multiple formats in listAll + // ═══════════════════════════════════════════════════════════════ + + public void testListAll_multipleFormats() throws IOException { + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput("_0.si", IOContext.DEFAULT)) { + out.writeString("lucene"); + } + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput("parquet/data.parquet", IOContext.DEFAULT)) { + out.writeString("parquet"); + } + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput("arrow/data.arrow", IOContext.DEFAULT)) { + out.writeString("arrow"); + } + + String[] files = dataFormatAwareStoreDirectory.listAll(); + List fileList = Arrays.asList(files); + assertTrue("Should contain lucene file", fileList.contains("_0.si")); + assertTrue("Should contain parquet file", fileList.contains("parquet/data.parquet")); + assertTrue("Should contain arrow file", fileList.contains("arrow/data.arrow")); + } + + // ═══════════════════════════════════════════════════════════════ + // sync with non-Lucene and multiple files + // ═══════════════════════════════════════════════════════════════ + + public void testSync_nonLuceneFile() throws IOException { + String fileIdentifier = "parquet/sync_test.parquet"; + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fileIdentifier, IOContext.DEFAULT)) { + out.writeString("sync test parquet"); + } + + // Should not throw + dataFormatAwareStoreDirectory.sync(Set.of(fileIdentifier)); + } + + public void testSync_multipleFiles() throws IOException { + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput("_sync1.si", IOContext.DEFAULT)) { + out.writeString("sync1"); + } + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput("parquet/sync2.parquet", IOContext.DEFAULT)) { + out.writeString("sync2"); + } + + // Should not throw with multiple files + dataFormatAwareStoreDirectory.sync(Set.of("_sync1.si", "parquet/sync2.parquet")); + } + + // ═══════════════════════════════════════════════════════════════ + // deleteFile with non-existent file + // ═══════════════════════════════════════════════════════════════ + + public void testDeleteFile_nonExistent_throws() { + expectThrows(IOException.class, () -> dataFormatAwareStoreDirectory.deleteFile("_nonexistent.si")); + } + + public void testDeleteFile_fileMetadata_nonExistent_throws() { + FileMetadata fm = new FileMetadata("parquet", "nonexistent.parquet"); + expectThrows(IOException.class, () -> dataFormatAwareStoreDirectory.deleteFile(fm.serialize())); + } + + // ═══════════════════════════════════════════════════════════════ + // toFileMetadata edge cases + // ═══════════════════════════════════════════════════════════════ + + public void testToFileMetadata_nestedPath() { + FileMetadata fm = DataFormatAwareStoreDirectory.toFileMetadata("custom/nested_file.data"); + assertEquals("custom", fm.dataFormat()); + assertEquals("nested_file.data", fm.file()); + } + + public void testToFileIdentifier_nullFormatTreatedAsDefault() { + FileMetadata fm = new FileMetadata(null, "_0.si"); + // null format should be treated as default (no prefix) + String identifier = DataFormatAwareStoreDirectory.toFileIdentifier(fm); + assertEquals("_0.si", identifier); + } + + public void testToFileIdentifier_emptyFormatTreatedAsDefault() { + FileMetadata fm = new FileMetadata("", "_0.si"); + String identifier = DataFormatAwareStoreDirectory.toFileIdentifier(fm); + assertEquals("_0.si", identifier); + } + + // ═══════════════════════════════════════════════════════════════ + // fileLength with serialized FileMetadata string + // ═══════════════════════════════════════════════════════════════ + + public void testFileLength_withSerializedString_lucene() throws IOException { + String fileName = "_len_serial.si"; + byte[] data = "serialized length test".getBytes(StandardCharsets.UTF_8); + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fileName, IOContext.DEFAULT)) { + out.writeBytes(data, data.length); + } + + // Access via serialized form + String serialized = new FileMetadata("lucene", fileName).serialize(); + assertEquals(data.length, dataFormatAwareStoreDirectory.fileLength(serialized)); + } + + // ═══════════════════════════════════════════════════════════════ + // Empty file checksum + // ═══════════════════════════════════════════════════════════════ + + public void testCalculateChecksum_emptyNonLuceneFile() throws IOException { + String fileIdentifier = "parquet/empty.parquet"; + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fileIdentifier, IOContext.DEFAULT)) { + // Write nothing - empty file + } + + long checksum = dataFormatAwareStoreDirectory.calculateChecksum(fileIdentifier); + + CRC32 crc32 = new CRC32(); + // CRC32 of empty data + assertEquals("CRC32 of empty file should match", crc32.getValue(), checksum); + } + + // ═══════════════════════════════════════════════════════════════ + // calculateUploadChecksum for lucene file + // ═══════════════════════════════════════════════════════════════ + + public void testCalculateUploadChecksum_lucene() throws IOException { + String fileName = "_upload_cksum_lucene.si"; + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fileName, IOContext.DEFAULT)) { + CodecUtil.writeHeader(out, "UploadTest", 1); + out.writeString("upload checksum lucene data"); + CodecUtil.writeFooter(out); + } + + FileMetadata fm = new FileMetadata("lucene", fileName); + String uploadChecksum = dataFormatAwareStoreDirectory.calculateUploadChecksum(fm.serialize()); + assertNotNull(uploadChecksum); + long parsedChecksum = Long.parseLong(uploadChecksum); + assertEquals(dataFormatAwareStoreDirectory.calculateChecksum(fm.serialize()), parsedChecksum); + } + + // ═══════════════════════════════════════════════════════════════ + // FileMetadata → "/" identifier → SubdirectoryAwareDirectory + // Path mapping & storage location tests + // ═══════════════════════════════════════════════════════════════ + + // --- Lucene files: no prefix, stored in /index/ --- + + public void testPathMapping_luceneFile_storedInIndexDir() throws IOException { + // Lucene files (no slash) should be stored in /index/ + FileMetadata fm = new FileMetadata("lucene", "_0.cfs"); + String identifier = DataFormatAwareStoreDirectory.toFileIdentifier(fm); + assertEquals("_0.cfs", identifier); // no prefix + assertFalse("Lucene identifier should not contain '/'", identifier.contains("/")); + + // Write and verify it's accessible + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(identifier, IOContext.DEFAULT)) { + out.writeString("lucene data"); + } + // Verify physical file exists in index directory + assertTrue(Files.exists(indexPath.resolve("_0.cfs"))); + } + + public void testPathMapping_segmentsFile_storedInIndexDir() throws IOException { + // segments_N files should be treated as lucene format (default) + FileMetadata fm = DataFormatAwareStoreDirectory.toFileMetadata("segments_1"); + assertEquals("lucene", fm.dataFormat()); + assertEquals("segments_1", fm.file()); + + String identifier = DataFormatAwareStoreDirectory.toFileIdentifier(fm); + assertEquals("segments_1", identifier); // no prefix + + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(identifier, IOContext.DEFAULT)) { + out.writeString("segments data"); + } + assertTrue("segments file should be in index dir", Files.exists(indexPath.resolve("segments_1"))); + } + + public void testPathMapping_segmentInfoFile_storedInIndexDir() throws IOException { + // _0.si (segment info) is a lucene file + FileMetadata fm = DataFormatAwareStoreDirectory.toFileMetadata("_0.si"); + assertEquals("lucene", fm.dataFormat()); + String identifier = DataFormatAwareStoreDirectory.toFileIdentifier(fm); + assertEquals("_0.si", identifier); + + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(identifier, IOContext.DEFAULT)) { + out.writeString("si data"); + } + assertTrue(Files.exists(indexPath.resolve("_0.si"))); + } + + // --- Metadata files: treated as default format, stored in /index/ --- + + public void testPathMapping_metadataFormat_storedInIndexDir() throws IOException { + // "metadata" is in INDEX_DIRECTORY_FORMATS, so no prefix is added + FileMetadata fm = new FileMetadata("metadata", "metadata__1__5__abc"); + String identifier = DataFormatAwareStoreDirectory.toFileIdentifier(fm); + assertEquals("metadata__1__5__abc", identifier); // no prefix + assertFalse("Metadata identifier should not contain '/'", identifier.contains("/")); + + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(identifier, IOContext.DEFAULT)) { + out.writeString("metadata content"); + } + assertTrue("Metadata file should be in index dir", Files.exists(indexPath.resolve("metadata__1__5__abc"))); + } + + public void testPathMapping_metadataFormat_getDataFormat() { + // A plain metadata filename (no slash) starting with "metadata" is treated as metadata format + assertEquals("metadata", dataFormatAwareStoreDirectory.getDataFormat("metadata__1__2__3")); + } + + // --- Parquet files: "parquet/" prefix, stored in /parquet/ --- + + public void testPathMapping_parquetFile_storedInSubdir() throws IOException { + FileMetadata fm = new FileMetadata("parquet", "_0_1.parquet"); + String identifier = DataFormatAwareStoreDirectory.toFileIdentifier(fm); + assertEquals("parquet/_0_1.parquet", identifier); + + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(identifier, IOContext.DEFAULT)) { + out.writeString("parquet data"); + } + assertTrue("Parquet file should be in parquet subdir", Files.exists(shardDataPath.resolve("parquet").resolve("_0_1.parquet"))); + assertFalse("Parquet file should NOT be in index dir", Files.exists(indexPath.resolve("_0_1.parquet"))); + } + + public void testPathMapping_parquetFile_fromIdentifier() { + FileMetadata fm = DataFormatAwareStoreDirectory.toFileMetadata("parquet/_0_1.parquet"); + assertEquals("parquet", fm.dataFormat()); + assertEquals("_0_1.parquet", fm.file()); + } + + // --- Arrow files: "arrow/" prefix, stored in /arrow/ --- + + public void testPathMapping_arrowFile_storedInSubdir() throws IOException { + FileMetadata fm = new FileMetadata("arrow", "data.arrow"); + String identifier = DataFormatAwareStoreDirectory.toFileIdentifier(fm); + assertEquals("arrow/data.arrow", identifier); + + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(identifier, IOContext.DEFAULT)) { + out.writeString("arrow data"); + } + assertTrue("Arrow file should be in arrow subdir", Files.exists(shardDataPath.resolve("arrow").resolve("data.arrow"))); + assertFalse("Arrow file should NOT be in index dir", Files.exists(indexPath.resolve("data.arrow"))); + } + + // --- Custom format: "custom/" prefix, stored in /custom/ --- + + public void testPathMapping_customFormat_storedInSubdir() throws IOException { + FileMetadata fm = new FileMetadata("custom", "myfile.dat"); + String identifier = DataFormatAwareStoreDirectory.toFileIdentifier(fm); + assertEquals("custom/myfile.dat", identifier); + + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(identifier, IOContext.DEFAULT)) { + out.writeString("custom data"); + } + assertTrue("Custom file should be in custom subdir", Files.exists(shardDataPath.resolve("custom").resolve("myfile.dat"))); + } + + // --- resolveFileName: serialized FileMetadata (with /) → native "/" identifier --- + + public void testResolveFileName_luceneSerialized() throws IOException { + // Write using plain name + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput("_0.si", IOContext.DEFAULT)) { + out.writeString("lucene data"); + } + + // Access using serialized FileMetadata — lucene serializes to plain name + String serialized = new FileMetadata("lucene", "_0.si").serialize(); + long length = dataFormatAwareStoreDirectory.fileLength(serialized); + assertTrue("Should resolve serialized lucene name", length > 0); + } + + public void testResolveFileName_parquetSerialized() throws IOException { + // Write using "/" identifier + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput("parquet/data.parquet", IOContext.DEFAULT)) { + out.writeString("parquet data"); + } + + // Access using serialized FileMetadata "parquet/data.parquet" + String serialized = new FileMetadata("parquet", "data.parquet").serialize(); + long length = dataFormatAwareStoreDirectory.fileLength(serialized); + assertTrue("Should resolve serialized parquet name", length > 0); + } + + public void testResolveFileName_metadataSerialized() throws IOException { + // Write using plain name (metadata is a default format) + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput("metadata__1__2__3", IOContext.DEFAULT)) { + out.writeString("metadata content"); + } + + // Access using serialized FileMetadata "metadata/metadata__1__2__3" + String serialized = new FileMetadata("metadata", "metadata__1__2__3").serialize(); + long length = dataFormatAwareStoreDirectory.fileLength(serialized); + assertTrue("Should resolve serialized metadata name", length > 0); + } + + // --- End-to-end: Write via FileMetadata, read via string identifier and vice versa --- + + public void testEndToEnd_writeViaFileMetadata_readViaString_lucene() throws IOException { + FileMetadata fm = new FileMetadata("lucene", "_e2e_lucene.si"); + byte[] data = "e2e lucene test".getBytes(StandardCharsets.UTF_8); + + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fm, IOContext.DEFAULT)) { + out.writeBytes(data, data.length); + } + + // Read using string identifier (no prefix for lucene) + try (IndexInput in = dataFormatAwareStoreDirectory.openInput("_e2e_lucene.si", IOContext.DEFAULT)) { + byte[] readData = new byte[data.length]; + in.readBytes(readData, 0, readData.length); + assertArrayEquals(data, readData); + } + } + + public void testEndToEnd_writeViaFileMetadata_readViaString_parquet() throws IOException { + FileMetadata fm = new FileMetadata("parquet", "e2e_data.parquet"); + byte[] data = "e2e parquet test".getBytes(StandardCharsets.UTF_8); + + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput(fm, IOContext.DEFAULT)) { + out.writeBytes(data, data.length); + } + + // Read using "/" string identifier + try (IndexInput in = dataFormatAwareStoreDirectory.openInput("parquet/e2e_data.parquet", IOContext.DEFAULT)) { + byte[] readData = new byte[data.length]; + in.readBytes(readData, 0, readData.length); + assertArrayEquals(data, readData); + } + } + + public void testEndToEnd_writeViaString_readViaFileMetadata_parquet() throws IOException { + byte[] data = "reverse e2e".getBytes(StandardCharsets.UTF_8); + + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput("parquet/rev_e2e.parquet", IOContext.DEFAULT)) { + out.writeBytes(data, data.length); + } + + // Read via FileMetadata + FileMetadata fm = new FileMetadata("parquet", "rev_e2e.parquet"); + try (IndexInput in = dataFormatAwareStoreDirectory.openInput(fm.serialize(), IOContext.DEFAULT)) { + byte[] readData = new byte[data.length]; + in.readBytes(readData, 0, readData.length); + assertArrayEquals(data, readData); + } + } + + // --- Physical path isolation: files of different formats don't collide --- + + public void testPathIsolation_sameFilenameInDifferentFormats() throws IOException { + byte[] luceneData = "lucene version".getBytes(StandardCharsets.UTF_8); + byte[] parquetData = "parquet version".getBytes(StandardCharsets.UTF_8); + byte[] arrowData = "arrow version".getBytes(StandardCharsets.UTF_8); + + // Write "data.file" in three different formats + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput("data.file", IOContext.DEFAULT)) { + out.writeBytes(luceneData, luceneData.length); + } + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput("parquet/data.file", IOContext.DEFAULT)) { + out.writeBytes(parquetData, parquetData.length); + } + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput("arrow/data.file", IOContext.DEFAULT)) { + out.writeBytes(arrowData, arrowData.length); + } + + // Verify different physical locations + assertTrue(Files.exists(indexPath.resolve("data.file"))); + assertTrue(Files.exists(shardDataPath.resolve("parquet").resolve("data.file"))); + assertTrue(Files.exists(shardDataPath.resolve("arrow").resolve("data.file"))); + + // Verify content is different (not overwritten) + assertEquals(luceneData.length, dataFormatAwareStoreDirectory.fileLength("data.file")); + assertEquals(parquetData.length, dataFormatAwareStoreDirectory.fileLength("parquet/data.file")); + assertEquals(arrowData.length, dataFormatAwareStoreDirectory.fileLength("arrow/data.file")); + } + + // --- Comprehensive toFileMetadata + toFileIdentifier round-trip for all formats --- + + public void testRoundtrip_allFormats() { + // Lucene + verifyRoundtrip("_0.cfs", "lucene", "_0.cfs"); + verifyRoundtrip("_0.si", "lucene", "_0.si"); + verifyRoundtrip("segments_1", "lucene", "segments_1"); + + // Non-lucene formats + verifyRoundtrip("parquet/_0_1.parquet", "parquet", "_0_1.parquet"); + verifyRoundtrip("arrow/data.arrow", "arrow", "data.arrow"); + verifyRoundtrip("custom/myfile.dat", "custom", "myfile.dat"); + } + + private void verifyRoundtrip(String identifier, String expectedFormat, String expectedFile) { + FileMetadata fm = DataFormatAwareStoreDirectory.toFileMetadata(identifier); + assertEquals("Format for " + identifier, expectedFormat, fm.dataFormat()); + assertEquals("File for " + identifier, expectedFile, fm.file()); + + String roundtripped = DataFormatAwareStoreDirectory.toFileIdentifier(fm); + assertEquals("Roundtrip for " + identifier, identifier, roundtripped); + } + + // --- isDefaultFormat edge cases (should not add prefix) --- + + public void testToFileIdentifier_defaultFormats_noPrefix() { + // "lucene" → no prefix + assertEquals("file.si", DataFormatAwareStoreDirectory.toFileIdentifier(new FileMetadata("lucene", "file.si"))); + // "LUCENE" (case-insensitive) → no prefix + assertEquals("file.si", DataFormatAwareStoreDirectory.toFileIdentifier(new FileMetadata("LUCENE", "file.si"))); + // "metadata" → no prefix + assertEquals("meta.dat", DataFormatAwareStoreDirectory.toFileIdentifier(new FileMetadata("metadata", "meta.dat"))); + // "METADATA" (case-insensitive) → no prefix + assertEquals("meta.dat", DataFormatAwareStoreDirectory.toFileIdentifier(new FileMetadata("METADATA", "meta.dat"))); + // null → no prefix + assertEquals("file.si", DataFormatAwareStoreDirectory.toFileIdentifier(new FileMetadata(null, "file.si"))); + // empty string → no prefix + assertEquals("file.si", DataFormatAwareStoreDirectory.toFileIdentifier(new FileMetadata("", "file.si"))); + } + + public void testToFileIdentifier_nonDefaultFormats_addPrefix() { + // Non-default formats always get "format/" prefix + assertEquals("parquet/data.parquet", DataFormatAwareStoreDirectory.toFileIdentifier(new FileMetadata("parquet", "data.parquet"))); + assertEquals("arrow/data.arrow", DataFormatAwareStoreDirectory.toFileIdentifier(new FileMetadata("arrow", "data.arrow"))); + assertEquals("orc/data.orc", DataFormatAwareStoreDirectory.toFileIdentifier(new FileMetadata("orc", "data.orc"))); + assertEquals("custom/my.file", DataFormatAwareStoreDirectory.toFileIdentifier(new FileMetadata("custom", "my.file"))); + } + + // --- listAll includes files from all formats with correct identifiers --- + + public void testListAll_returnsCorrectIdentifiers() throws IOException { + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput("_0.si", IOContext.DEFAULT)) { + out.writeString("lucene"); + } + try (IndexOutput out = dataFormatAwareStoreDirectory.createOutput("parquet/data.parquet", IOContext.DEFAULT)) { + out.writeString("parquet"); + } + + String[] files = dataFormatAwareStoreDirectory.listAll(); + List fileList = Arrays.asList(files); + + // Lucene files should appear as plain names (no prefix) + assertTrue("Lucene file listed as plain name", fileList.contains("_0.si")); + assertFalse("Lucene file should NOT have lucene/ prefix", fileList.contains("lucene/_0.si")); + + // Non-lucene files should appear with serialized "format/file" form + assertTrue("Parquet file listed with serialized form", fileList.contains("parquet/data.parquet")); + } + + // --- getDataFormat comprehensive --- + + public void testGetDataFormat_comprehensive() { + // Plain filenames → "lucene" + assertEquals("lucene", dataFormatAwareStoreDirectory.getDataFormat("_0.si")); + assertEquals("lucene", dataFormatAwareStoreDirectory.getDataFormat("_0.cfs")); + assertEquals("lucene", dataFormatAwareStoreDirectory.getDataFormat("_0.cfe")); + assertEquals("lucene", dataFormatAwareStoreDirectory.getDataFormat("segments_1")); + assertEquals("lucene", dataFormatAwareStoreDirectory.getDataFormat("write.lock")); + + // Prefixed filenames → format name + assertEquals("parquet", dataFormatAwareStoreDirectory.getDataFormat("parquet/data.parquet")); + assertEquals("arrow", dataFormatAwareStoreDirectory.getDataFormat("arrow/data.arrow")); + assertEquals("orc", dataFormatAwareStoreDirectory.getDataFormat("orc/data.orc")); + assertEquals("custom", dataFormatAwareStoreDirectory.getDataFormat("custom/myfile.dat")); + } +} diff --git a/server/src/test/java/org/opensearch/index/store/DefaultDataFormatAwareStoreDirectoryFactoryTests.java b/server/src/test/java/org/opensearch/index/store/DefaultDataFormatAwareStoreDirectoryFactoryTests.java new file mode 100644 index 0000000000000..d1a47e9710661 --- /dev/null +++ b/server/src/test/java/org/opensearch/index/store/DefaultDataFormatAwareStoreDirectoryFactoryTests.java @@ -0,0 +1,208 @@ +/* + * 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.store; + +import org.apache.lucene.store.FSDirectory; +import org.opensearch.cluster.metadata.IndexMetadata; +import org.opensearch.common.settings.Settings; +import org.opensearch.core.index.Index; +import org.opensearch.core.index.shard.ShardId; +import org.opensearch.index.IndexSettings; +import org.opensearch.index.engine.dataformat.DataFormatPlugin; +import org.opensearch.index.engine.dataformat.DataFormatRegistry; +import org.opensearch.index.shard.ShardPath; +import org.opensearch.plugins.IndexStorePlugin; +import org.opensearch.plugins.PluginsService; +import org.opensearch.plugins.SearchBackEndPlugin; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; + +import static org.opensearch.cluster.metadata.IndexMetadata.SETTING_INDEX_UUID; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** + * Unit tests for {@link DefaultDataFormatAwareStoreDirectoryFactory}. + */ +public class DefaultDataFormatAwareStoreDirectoryFactoryTests extends OpenSearchTestCase { + + private ShardPath createShardPath(Path tempDir) throws IOException { + String indexUUID = "test-index-uuid"; + int shardId = 0; + Path shardDataPath = tempDir.resolve(indexUUID).resolve(Integer.toString(shardId)); + Path indexPath = shardDataPath.resolve(ShardPath.INDEX_FOLDER_NAME); + Files.createDirectories(indexPath); + + ShardId sid = new ShardId(new Index("test-index", indexUUID), shardId); + return new ShardPath(false, shardDataPath, shardDataPath, sid); + } + + private IndexSettings createIndexSettings() { + Settings settings = Settings.builder() + .put(IndexMetadata.SETTING_VERSION_CREATED, org.opensearch.Version.CURRENT) + .put(SETTING_INDEX_UUID, "test-index-uuid") + .build(); + IndexMetadata metadata = IndexMetadata.builder("test-index").settings(settings).numberOfShards(1).numberOfReplicas(0).build(); + return new IndexSettings(metadata, Settings.EMPTY); + } + + private DataFormatRegistry createEmptyDataFormatRegistry() { + PluginsService pluginsService = mock(PluginsService.class); + when(pluginsService.filterPlugins(DataFormatPlugin.class)).thenReturn(List.of()); + when(pluginsService.filterPlugins(SearchBackEndPlugin.class)).thenReturn(List.of()); + return new DataFormatRegistry(pluginsService); + } + + private IndexStorePlugin.DirectoryFactory createFsDirectoryFactory() { + return new IndexStorePlugin.DirectoryFactory() { + @Override + public org.apache.lucene.store.Directory newDirectory(IndexSettings indexSettings, ShardPath shardPath) throws IOException { + return FSDirectory.open(shardPath.resolveIndex()); + } + + @Override + public org.apache.lucene.store.Directory newFSDirectory( + Path location, + org.apache.lucene.store.LockFactory lockFactory, + IndexSettings indexSettings + ) throws IOException { + return FSDirectory.open(location, lockFactory); + } + }; + } + + // ═══════════════════════════════════════════════════════════════ + // newDataFormatAwareStoreDirectory Tests + // ═══════════════════════════════════════════════════════════════ + + public void testNewDataFormatAwareStoreDirectory_CreatesSuccessfully() throws IOException { + DataFormatRegistry registry = createEmptyDataFormatRegistry(); + DefaultDataFormatAwareStoreDirectoryFactory factory = new DefaultDataFormatAwareStoreDirectoryFactory(); + Path tempDir = createTempDir(); + ShardPath shardPath = createShardPath(tempDir); + IndexSettings indexSettings = createIndexSettings(); + + DataFormatAwareStoreDirectory directory = factory.newDataFormatAwareStoreDirectory( + indexSettings, + shardPath.getShardId(), + shardPath, + createFsDirectoryFactory(), + registry + ); + + assertNotNull("Factory should create a non-null DataFormatAwareStoreDirectory", directory); + } + + public void testNewDataFormatAwareStoreDirectory_HasCorrectShardPath() throws IOException { + DataFormatRegistry registry = createEmptyDataFormatRegistry(); + DefaultDataFormatAwareStoreDirectoryFactory factory = new DefaultDataFormatAwareStoreDirectoryFactory(); + Path tempDir = createTempDir(); + ShardPath shardPath = createShardPath(tempDir); + IndexSettings indexSettings = createIndexSettings(); + + DataFormatAwareStoreDirectory directory = factory.newDataFormatAwareStoreDirectory( + indexSettings, + shardPath.getShardId(), + shardPath, + createFsDirectoryFactory(), + registry + ); + + assertEquals(shardPath, directory.getShardPath()); + } + + public void testNewDataFormatAwareStoreDirectory_CanListFiles() throws IOException { + DataFormatRegistry registry = createEmptyDataFormatRegistry(); + DefaultDataFormatAwareStoreDirectoryFactory factory = new DefaultDataFormatAwareStoreDirectoryFactory(); + Path tempDir = createTempDir(); + ShardPath shardPath = createShardPath(tempDir); + IndexSettings indexSettings = createIndexSettings(); + + DataFormatAwareStoreDirectory directory = factory.newDataFormatAwareStoreDirectory( + indexSettings, + shardPath.getShardId(), + shardPath, + createFsDirectoryFactory(), + registry + ); + + // Should not throw + String[] files = directory.listAll(); + assertNotNull(files); + } + + public void testNewDataFormatAwareStoreDirectory_MultipleCalls_CreatesSeparateInstances() throws IOException { + DataFormatRegistry registry = createEmptyDataFormatRegistry(); + DefaultDataFormatAwareStoreDirectoryFactory factory = new DefaultDataFormatAwareStoreDirectoryFactory(); + Path tempDir1 = createTempDir(); + Path tempDir2 = createTempDir(); + ShardPath shardPath1 = createShardPath(tempDir1); + ShardPath shardPath2 = createShardPath(tempDir2); + IndexSettings indexSettings = createIndexSettings(); + + DataFormatAwareStoreDirectory dir1 = factory.newDataFormatAwareStoreDirectory( + indexSettings, + shardPath1.getShardId(), + shardPath1, + createFsDirectoryFactory(), + registry + ); + DataFormatAwareStoreDirectory dir2 = factory.newDataFormatAwareStoreDirectory( + indexSettings, + shardPath2.getShardId(), + shardPath2, + createFsDirectoryFactory(), + registry + ); + + assertNotNull(dir1); + assertNotNull(dir2); + assertNotSame("Each call should create a new instance", dir1, dir2); + } + + public void testNewDataFormatAwareStoreDirectory_InvalidPath_ThrowsIOException() throws IOException { + DataFormatRegistry registry = createEmptyDataFormatRegistry(); + DefaultDataFormatAwareStoreDirectoryFactory factory = new DefaultDataFormatAwareStoreDirectoryFactory(); + IndexSettings indexSettings = createIndexSettings(); + + // Create a valid shard path structure (must end with shardId, parent with indexUUID) + // but place a regular file where the "index" directory should be, so FSDirectory.open() fails + String indexUUID = "test-index-uuid"; + int shardId = 0; + Path tempDir = createTempDir(); + Path shardDataPath = tempDir.resolve(indexUUID).resolve(Integer.toString(shardId)); + Files.createDirectories(shardDataPath); + // Create a FILE named "index" instead of a directory — FSDirectory.open() will fail + Path indexFile = shardDataPath.resolve(ShardPath.INDEX_FOLDER_NAME); + Files.createFile(indexFile); + + ShardId sid = new ShardId(new Index("test-index", indexUUID), shardId); + ShardPath invalidShardPath = new ShardPath(false, shardDataPath, shardDataPath, sid); + + // This should trigger the catch block which wraps the exception as IOException + IOException exception = expectThrows( + IOException.class, + () -> factory.newDataFormatAwareStoreDirectory( + indexSettings, + invalidShardPath.getShardId(), + invalidShardPath, + createFsDirectoryFactory(), + registry + ) + ); + assertTrue( + "Exception message should mention shard, but was: " + exception.getMessage(), + exception.getMessage().contains("Failed to create DataFormatAwareStoreDirectory for shard") + ); + } +} diff --git a/server/src/test/java/org/opensearch/index/store/FileMetadataTests.java b/server/src/test/java/org/opensearch/index/store/FileMetadataTests.java new file mode 100644 index 0000000000000..208fde4527c54 --- /dev/null +++ b/server/src/test/java/org/opensearch/index/store/FileMetadataTests.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.index.store; + +import org.opensearch.test.OpenSearchTestCase; + +public class FileMetadataTests extends OpenSearchTestCase { + + // ═══════════════════════════════════════════════════════════════ + // Two-arg constructor tests + // ═══════════════════════════════════════════════════════════════ + + public void testConstructorTwoArgs() { + FileMetadata fm = new FileMetadata("parquet", "_0_1.parquet"); + assertEquals("parquet", fm.dataFormat()); + assertEquals("_0_1.parquet", fm.file()); + } + + public void testConstructorTwoArgs_lucene() { + FileMetadata fm = new FileMetadata("lucene", "_0.si"); + assertEquals("lucene", fm.dataFormat()); + assertEquals("_0.si", fm.file()); + } + + // ═══════════════════════════════════════════════════════════════ + // Single-arg constructor tests + // ═══════════════════════════════════════════════════════════════ + + public void testConstructorSingleArg_luceneFileWithoutDelimiter() { + // A plain filename without delimiter defaults to "lucene" + FileMetadata fm = new FileMetadata("_0.si"); + assertEquals("lucene", fm.dataFormat()); + assertEquals("_0.si", fm.file()); + } + + public void testConstructorSingleArg_withDelimiter() { + FileMetadata fm = new FileMetadata("parquet/_0_1.parquet"); + assertEquals("parquet", fm.dataFormat()); + assertEquals("_0_1.parquet", fm.file()); + } + + public void testConstructorSingleArg_withDelimiterLucene() { + // Plain lucene files don't have a prefix, so single-arg defaults to "lucene" + FileMetadata fm = new FileMetadata("_0.si"); + assertEquals("lucene", fm.dataFormat()); + assertEquals("_0.si", fm.file()); + } + + public void testConstructorSingleArg_metadataKey() { + // Files starting with "metadata" and not containing delimiter are treated as metadata format + FileMetadata fm = new FileMetadata("metadata__1__2__3"); + assertEquals("metadata", fm.dataFormat()); + assertEquals("metadata__1__2__3", fm.file()); + } + + public void testConstructorSingleArg_metadataKeyWithDelimiter() { + // If it contains delimiter, parse normally even if it starts with "metadata" + FileMetadata fm = new FileMetadata("metadata/metadata__1__2__3"); + assertEquals("metadata", fm.dataFormat()); + assertEquals("metadata__1__2__3", fm.file()); + } + + // ═══════════════════════════════════════════════════════════════ + // serialize / toString + // ═══════════════════════════════════════════════════════════════ + + public void testSerialize() { + FileMetadata fm = new FileMetadata("parquet", "_0_1.parquet"); + assertEquals("parquet/_0_1.parquet", fm.serialize()); + } + + public void testSerialize_lucene() { + FileMetadata fm = new FileMetadata("lucene", "_0.si"); + assertEquals("_0.si", fm.serialize()); + } + + public void testToStringEqualsSerialized() { + FileMetadata fm = new FileMetadata("arrow", "data.arrow"); + assertEquals(fm.serialize(), fm.toString()); + } + + // ═══════════════════════════════════════════════════════════════ + // Roundtrip: two-arg -> serialize -> single-arg + // ═══════════════════════════════════════════════════════════════ + + public void testRoundtrip() { + FileMetadata original = new FileMetadata("parquet", "_0_1.parquet"); + String serialized = original.serialize(); + FileMetadata deserialized = new FileMetadata(serialized); + assertEquals(original, deserialized); + assertEquals(original.file(), deserialized.file()); + assertEquals(original.dataFormat(), deserialized.dataFormat()); + } + + public void testRoundtrip_lucene() { + FileMetadata original = new FileMetadata("lucene", "_0.cfs"); + String serialized = original.serialize(); + FileMetadata deserialized = new FileMetadata(serialized); + assertEquals(original, deserialized); + } + + // ═══════════════════════════════════════════════════════════════ + // equals / hashCode + // ═══════════════════════════════════════════════════════════════ + + public void testEquals_sameValues() { + FileMetadata fm1 = new FileMetadata("parquet", "_0.parquet"); + FileMetadata fm2 = new FileMetadata("parquet", "_0.parquet"); + assertEquals(fm1, fm2); + assertEquals(fm1.hashCode(), fm2.hashCode()); + } + + public void testNotEqual_differentFile() { + FileMetadata fm1 = new FileMetadata("parquet", "_0.parquet"); + FileMetadata fm2 = new FileMetadata("parquet", "_1.parquet"); + assertNotEquals(fm1, fm2); + } + + public void testNotEqual_differentFormat() { + FileMetadata fm1 = new FileMetadata("parquet", "_0.parquet"); + FileMetadata fm2 = new FileMetadata("arrow", "_0.parquet"); + assertNotEquals(fm1, fm2); + } + + public void testEquals_null() { + FileMetadata fm = new FileMetadata("lucene", "_0.si"); + assertNotEquals(null, fm); + } + + public void testEquals_differentType() { + FileMetadata fm = new FileMetadata("lucene", "_0.si"); + assertNotEquals("_0.si", fm); + } + + public void testEquals_self() { + FileMetadata fm = new FileMetadata("lucene", "_0.si"); + assertEquals(fm, fm); + } + + // ═══════════════════════════════════════════════════════════════ + // DELIMITER constant + // ═══════════════════════════════════════════════════════════════ + + public void testDelimiterConstant() { + assertEquals("/", FileMetadata.DELIMITER); + } + + // ═══════════════════════════════════════════════════════════════ + // Edge cases - single-arg constructor + // ═══════════════════════════════════════════════════════════════ + + public void testConstructorSingleArg_segmentsFile() { + // Segments files (e.g., "segments_1") should default to "lucene" format + FileMetadata fm = new FileMetadata("segments_1"); + assertEquals("lucene", fm.dataFormat()); + assertEquals("segments_1", fm.file()); + } + + public void testConstructorSingleArg_metadataExactString() { + // The exact string "metadata" (without underscores/suffixes) should still be treated as metadata + FileMetadata fm = new FileMetadata("metadata"); + assertEquals("metadata", fm.dataFormat()); + assertEquals("metadata", fm.file()); + } + + public void testConstructorSingleArg_arrowWithDelimiter() { + FileMetadata fm = new FileMetadata("arrow/data.arrow"); + assertEquals("arrow", fm.dataFormat()); + assertEquals("data.arrow", fm.file()); + } + + public void testConstructorSingleArg_customFormat() { + FileMetadata fm = new FileMetadata("myformat/data.custom"); + assertEquals("myformat", fm.dataFormat()); + assertEquals("data.custom", fm.file()); + } + + // ═══════════════════════════════════════════════════════════════ + // Roundtrip - metadata format + // ═══════════════════════════════════════════════════════════════ + + public void testRoundtrip_metadata() { + FileMetadata original = new FileMetadata("metadata", "metadata__1__2__3"); + String serialized = original.serialize(); + assertEquals("metadata/metadata__1__2__3", serialized); + FileMetadata deserialized = new FileMetadata(serialized); + assertEquals(original, deserialized); + } + + public void testRoundtrip_arrow() { + FileMetadata original = new FileMetadata("arrow", "data.arrow"); + String serialized = original.serialize(); + FileMetadata deserialized = new FileMetadata(serialized); + assertEquals(original, deserialized); + } + + // ═══════════════════════════════════════════════════════════════ + // hashCode - inequality + // ═══════════════════════════════════════════════════════════════ + + public void testHashCode_differentObjects_likelyDifferent() { + FileMetadata fm1 = new FileMetadata("parquet", "_0.parquet"); + FileMetadata fm2 = new FileMetadata("parquet", "_1.parquet"); + // Different content should (very likely) produce different hashCodes + assertNotEquals(fm1.hashCode(), fm2.hashCode()); + } + + public void testHashCode_differentFormat_likelyDifferent() { + FileMetadata fm1 = new FileMetadata("parquet", "_0.data"); + FileMetadata fm2 = new FileMetadata("arrow", "_0.data"); + assertNotEquals(fm1.hashCode(), fm2.hashCode()); + } + + // ═══════════════════════════════════════════════════════════════ + // Use as Map key + // ═══════════════════════════════════════════════════════════════ + + public void testUsableAsMapKey() { + java.util.Map map = new java.util.HashMap<>(); + FileMetadata key1 = new FileMetadata("parquet", "_0.parquet"); + FileMetadata key2 = new FileMetadata("parquet", "_0.parquet"); // equal to key1 + + map.put(key1, "value1"); + assertEquals("value1", map.get(key2)); // should find by equal key + assertEquals(1, map.size()); + } + + public void testUsableAsMapKey_differentFormats() { + java.util.Map map = new java.util.HashMap<>(); + FileMetadata key1 = new FileMetadata("parquet", "_0.data"); + FileMetadata key2 = new FileMetadata("arrow", "_0.data"); + + map.put(key1, "parquet-value"); + map.put(key2, "arrow-value"); + assertEquals(2, map.size()); + assertEquals("parquet-value", map.get(key1)); + assertEquals("arrow-value", map.get(key2)); + } + + // ═══════════════════════════════════════════════════════════════ + // Constructor consistency: two-arg vs single-arg from serialized + // ═══════════════════════════════════════════════════════════════ + + public void testTwoArgAndSingleArgConsistency() { + FileMetadata twoArg = new FileMetadata("parquet", "_0.parquet"); + FileMetadata singleArg = new FileMetadata(twoArg.serialize()); + assertEquals(twoArg, singleArg); + assertEquals(twoArg.file(), singleArg.file()); + assertEquals(twoArg.dataFormat(), singleArg.dataFormat()); + assertEquals(twoArg.hashCode(), singleArg.hashCode()); + } + + public void testFileAndDataFormatImmutable() { + FileMetadata fm = new FileMetadata("lucene", "_0.si"); + // Verify accessors return consistent values on repeated calls + assertEquals("lucene", fm.dataFormat()); + assertEquals("lucene", fm.dataFormat()); + assertEquals("_0.si", fm.file()); + assertEquals("_0.si", fm.file()); + } +} diff --git a/server/src/test/java/org/opensearch/index/store/RemoteSegmentStoreDirectoryTests.java b/server/src/test/java/org/opensearch/index/store/RemoteSegmentStoreDirectoryTests.java index a30074f35726d..6d54407abc01c 100644 --- a/server/src/test/java/org/opensearch/index/store/RemoteSegmentStoreDirectoryTests.java +++ b/server/src/test/java/org/opensearch/index/store/RemoteSegmentStoreDirectoryTests.java @@ -28,16 +28,20 @@ import org.opensearch.common.io.VersionedCodecStreamWrapper; import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.common.lucene.store.ByteArrayIndexInput; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Settings; import org.opensearch.core.action.ActionListener; import org.opensearch.core.common.bytes.BytesReference; import org.opensearch.core.index.Index; import org.opensearch.core.index.shard.ShardId; +import org.opensearch.index.engine.exec.coord.SegmentInfosCatalogSnapshot; import org.opensearch.index.remote.RemoteStoreEnums.PathHashAlgorithm; import org.opensearch.index.remote.RemoteStoreEnums.PathType; import org.opensearch.index.remote.RemoteStorePathStrategy; import org.opensearch.index.remote.RemoteStoreUtils; import org.opensearch.index.store.remote.metadata.RemoteSegmentMetadata; import org.opensearch.index.store.remote.metadata.RemoteSegmentMetadataHandlerFactory; +import org.opensearch.indices.RemoteStoreSettings; import org.opensearch.test.MockLogAppender; import org.opensearch.test.junit.annotations.TestLogging; import org.opensearch.threadpool.ThreadPool; @@ -64,9 +68,10 @@ import static org.opensearch.test.RemoteStoreTestUtils.createMetadataFileBytes; import static org.opensearch.test.RemoteStoreTestUtils.getDummyMetadata; import static org.hamcrest.CoreMatchers.is; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.argThat; -import static org.mockito.Mockito.any; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.eq; @@ -81,6 +86,10 @@ public class RemoteSegmentStoreDirectoryTests extends BaseRemoteSegmentStoreDire @Before public void setup() throws IOException { + new RemoteStoreSettings( + Settings.builder().put(RemoteStoreSettings.CLUSTER_REMOTE_STORE_PINNED_TIMESTAMP_ENABLED.getKey(), false).build(), + new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS) + ); setupRemoteSegmentStoreDirectory(); } @@ -129,6 +138,39 @@ public void testUploadedSegmentMetadataFromStringException() { ); } + public void testUploadedSegmentMetadataFromStringWithFormatSlash() { + // Format-aware originalFilename uses "/" delimiter: "parquet/_0.pqt" + String input = "parquet/_0.pqt::_0.pqt__uuidxyz::4567::372000::" + Version.LATEST.major; + RemoteSegmentStoreDirectory.UploadedSegmentMetadata metadata = RemoteSegmentStoreDirectory.UploadedSegmentMetadata.fromString( + input + ); + assertEquals("parquet/_0.pqt", metadata.getOriginalFilename()); + assertEquals("_0.pqt__uuidxyz", metadata.getUploadedFilename()); + assertEquals("4567", metadata.getChecksum()); + assertEquals(372000L, metadata.getLength()); + assertEquals(input, metadata.toString()); + } + + public void testUploadedSegmentMetadataRoundTripWithFormatSlash() { + // Create metadata with format-aware originalFilename, serialize, deserialize + RemoteSegmentStoreDirectory.UploadedSegmentMetadata metadata = new RemoteSegmentStoreDirectory.UploadedSegmentMetadata( + "parquet/_0.pqt", + "_0.pqt__uuid123", + "9999", + 5000 + ); + metadata.setWrittenByMajor(Version.LATEST.major); + String serialized = metadata.toString(); + RemoteSegmentStoreDirectory.UploadedSegmentMetadata deserialized = RemoteSegmentStoreDirectory.UploadedSegmentMetadata.fromString( + serialized + ); + assertEquals(metadata.getOriginalFilename(), deserialized.getOriginalFilename()); + assertEquals(metadata.getUploadedFilename(), deserialized.getUploadedFilename()); + assertEquals(metadata.getChecksum(), deserialized.getChecksum()); + assertEquals(metadata.getLength(), deserialized.getLength()); + assertEquals(serialized, deserialized.toString()); + } + public void testGetPrimaryTermGenerationUuid() { String[] filenameTokens = "abc__9223372036854775795__9223372036854775784__uuid_xyz".split(SEPARATOR); assertEquals(12, RemoteSegmentStoreDirectory.MetadataFilenameUtils.getPrimaryTerm(filenameTokens)); @@ -214,7 +256,7 @@ public void testDeleteFileException() throws IOException { populateMetadata(); remoteSegmentStoreDirectory.init(); - doThrow(new IOException("Error")).when(remoteDataDirectory).deleteFile(any()); + doThrow(new IOException("Error")).when(remoteDataDirectory).deleteFile(anyString()); assertThrows(IOException.class, () -> remoteSegmentStoreDirectory.deleteFile("_0.si")); } @@ -259,6 +301,7 @@ public void testOpenInput() throws IOException { remoteSegmentStoreDirectory.init(); IndexInput indexInput = mock(IndexInput.class); + // Mock String-based openInput when(remoteDataDirectory.openInput(startsWith("_0.si"), anyLong(), eq(IOContext.DEFAULT))).thenReturn(indexInput); assertEquals(indexInput, remoteSegmentStoreDirectory.openInput("_0.si", IOContext.DEFAULT)); @@ -272,6 +315,7 @@ public void testOpenInputException() throws IOException { populateMetadata(); remoteSegmentStoreDirectory.init(); + // Mock String-based openInput to throw when(remoteDataDirectory.openInput(startsWith("_0.si"), anyLong(), eq(IOContext.DEFAULT))).thenThrow(new IOException("Error")); assertThrows(IOException.class, () -> remoteSegmentStoreDirectory.openInput("_0.si", IOContext.DEFAULT)); @@ -366,7 +410,9 @@ public void testIsAcquiredException() throws IOException { private List getDummyMetadataFiles(int count) { List sortedMetadataFiles = new ArrayList<>(); for (int counter = 0; counter < count; counter++) { - sortedMetadataFiles.add(RemoteSegmentStoreDirectory.MetadataFilenameUtils.getMetadataFilename(counter, 23, 34, 1, 1, "node-1")); + sortedMetadataFiles.add( + RemoteSegmentStoreDirectory.MetadataFilenameUtils.getMetadataFilename(counter, 23, 34, 1, 1, "node-1", 0L) + ); } return sortedMetadataFiles; } @@ -1022,7 +1068,7 @@ public void testDeleteStaleCommitsNoDeletesDueToLocks() throws Exception { remoteSegmentStoreDirectory.deleteStaleSegmentsAsync(1); assertBusy(() -> assertThat(remoteSegmentStoreDirectory.canDeleteStaleCommits.get(), is(true))); - verify(remoteMetadataDirectory, times(0)).deleteFile(any()); + verify(remoteMetadataDirectory, times(0)).deleteFile(anyString()); } public void testDeleteStaleCommitsExceptionWhileFetchingLocks() throws Exception { @@ -1035,7 +1081,7 @@ public void testDeleteStaleCommitsExceptionWhileFetchingLocks() throws Exception // We are passing lastNMetadataFilesToKeep=2 here so that oldest 1 metadata file will be deleted remoteSegmentStoreDirectory.deleteStaleSegmentsAsync(1); - verify(remoteMetadataDirectory, times(0)).deleteFile(any()); + verify(remoteMetadataDirectory, times(0)).deleteFile(anyString()); } public void testDeleteStaleCommitsDeleteDedup() throws Exception { @@ -1345,6 +1391,167 @@ public void testInitializeToSpecificTimestampMatchingMdFile() throws IOException assertTrue(uploadedSegments.containsKey("_0.cfs")); } + // ═══════════════════════════════════════════════════════════════ + // Tests for new CatalogSnapshot-based uploadMetadata + // ═══════════════════════════════════════════════════════════════ + + public void testUploadMetadataWithCatalogSnapshot() throws IOException { + indexDocs(142364, 5); + flushShard(indexShard, true); + SegmentInfos segInfos = indexShard.store().readLastCommittedSegmentsInfo(); + long primaryTerm = indexShard.getLatestReplicationCheckpoint().getPrimaryTerm(); + String primaryTermLong = RemoteStoreUtils.invertLong(primaryTerm); + long generation = segInfos.getGeneration(); + String generationLong = RemoteStoreUtils.invertLong(generation); + String latestMetadataFileName = "metadata__" + primaryTermLong + "__" + generationLong + "__abc"; + List metadataFiles = List.of(latestMetadataFileName); + when( + remoteMetadataDirectory.listFilesByPrefixInLexicographicOrder( + RemoteSegmentStoreDirectory.MetadataFilenameUtils.METADATA_PREFIX, + METADATA_FILES_TO_FETCH + ) + ).thenReturn(metadataFiles); + Map> metadataFilenameContentMapping = Map.of( + latestMetadataFileName, + getDummyMetadata("_0", (int) generation) + ); + when(remoteMetadataDirectory.getBlobStream(latestMetadataFileName)).thenReturn( + createMetadataFileBytes( + metadataFilenameContentMapping.get(latestMetadataFileName), + indexShard.getLatestReplicationCheckpoint(), + segmentInfos + ) + ); + + remoteSegmentStoreDirectory.init(); + + Directory storeDirectory = mock(Directory.class); + BytesStreamOutput output = new BytesStreamOutput(); + IndexOutput indexOutput = new OutputStreamIndexOutput("segment metadata", "metadata output stream", output, 4096); + when(storeDirectory.createOutput(startsWith("metadata__" + primaryTermLong + "__" + generationLong), eq(IOContext.DEFAULT))) + .thenReturn(indexOutput); + + // Create CatalogSnapshot from SegmentInfos + SegmentInfosCatalogSnapshot catalogSnapshot = new SegmentInfosCatalogSnapshot(segInfos); + + remoteSegmentStoreDirectory.uploadMetadata( + segInfos.files(true), + catalogSnapshot, + storeDirectory, + generation, + indexShard.getLatestReplicationCheckpoint(), + "" + ); + + verify(remoteMetadataDirectory).copyFrom( + eq(storeDirectory), + startsWith("metadata__" + primaryTermLong + "__" + generationLong), + startsWith("metadata__" + primaryTermLong + "__" + generationLong), + eq(IOContext.DEFAULT) + ); + } + + public void testUploadMetadataWithCatalogSnapshot_MissingSegment() throws IOException { + populateMetadata(); + remoteSegmentStoreDirectory.init(); + + Directory storeDirectory = mock(Directory.class); + IndexOutput indexOutput = mock(IndexOutput.class); + + String generation = RemoteStoreUtils.invertLong(segmentInfos.getGeneration()); + long primaryTermLong = indexShard.getLatestReplicationCheckpoint().getPrimaryTerm(); + String primaryTerm = RemoteStoreUtils.invertLong(primaryTermLong); + when(storeDirectory.createOutput(startsWith("metadata__" + primaryTerm + "__" + generation), eq(IOContext.DEFAULT))).thenReturn( + indexOutput + ); + + SegmentInfosCatalogSnapshot catalogSnapshot = new SegmentInfosCatalogSnapshot(segmentInfos); + + Collection segmentFiles = List.of("_123.si"); + assertThrows( + NoSuchFileException.class, + () -> remoteSegmentStoreDirectory.uploadMetadata( + segmentFiles, + catalogSnapshot, + storeDirectory, + 12L, + indexShard.getLatestReplicationCheckpoint(), + "" + ) + ); + verify(indexOutput).close(); + } + + // ═══════════════════════════════════════════════════════════════ + // Tests for UploadedSegmentMetadata with format-aware filenames + // ═══════════════════════════════════════════════════════════════ + + public void testUploadedSegmentMetadataFromString_WithFormatDelimiter() { + // Format: originalFilename::uploadedFilename::checksum::length::writtenByMajor + // Where originalFilename uses "/" convention (e.g., "parquet/_0.parquet") + String metadataString = "parquet/_0.parquet::_0.parquet__UUID1::checksum456::200::" + Version.LATEST.major; + RemoteSegmentStoreDirectory.UploadedSegmentMetadata metadata = RemoteSegmentStoreDirectory.UploadedSegmentMetadata.fromString( + metadataString + ); + + assertEquals("parquet/_0.parquet", metadata.getOriginalFilename()); + assertEquals("_0.parquet__UUID1", metadata.getUploadedFilename()); + assertEquals("checksum456", metadata.getChecksum()); + assertEquals(200, metadata.getLength()); + } + + public void testUploadedSegmentMetadataToString_WithFormatDelimiter() { + RemoteSegmentStoreDirectory.UploadedSegmentMetadata metadata = new RemoteSegmentStoreDirectory.UploadedSegmentMetadata( + "parquet/_0.parquet", + "_0.parquet__UUID1", + "checksum456", + 200 + ); + metadata.setWrittenByMajor(Version.LATEST.major); + + String result = metadata.toString(); + assertTrue("toString should contain parquet/", result.contains("parquet/_0.parquet")); + assertTrue("toString should contain uploaded filename", result.contains("_0.parquet__UUID1")); + + // Verify round-trip + RemoteSegmentStoreDirectory.UploadedSegmentMetadata parsed = RemoteSegmentStoreDirectory.UploadedSegmentMetadata.fromString(result); + assertEquals("parquet/_0.parquet", parsed.getOriginalFilename()); + assertEquals("_0.parquet__UUID1", parsed.getUploadedFilename()); + } + + // ═══════════════════════════════════════════════════════════════ + // Tests for readLatestNMetadataFiles + // ═══════════════════════════════════════════════════════════════ + + public void testReadLatestNMetadataFiles_Empty() throws IOException { + when( + remoteMetadataDirectory.listFilesByPrefixInLexicographicOrder( + RemoteSegmentStoreDirectory.MetadataFilenameUtils.METADATA_PREFIX, + 3 + ) + ).thenReturn(new ArrayList<>()); + + Map result = remoteSegmentStoreDirectory.readLatestNMetadataFiles(3); + assertNotNull(result); + assertEquals(0, result.size()); + } + + public void testReadLatestNMetadataFiles_SingleFile() throws IOException { + populateMetadata(); + + when( + remoteMetadataDirectory.listFilesByPrefixInLexicographicOrder( + RemoteSegmentStoreDirectory.MetadataFilenameUtils.METADATA_PREFIX, + 1 + ) + ).thenReturn(List.of(metadataFilename)); + + Map result = remoteSegmentStoreDirectory.readLatestNMetadataFiles(1); + assertNotNull(result); + assertEquals(1, result.size()); + assertTrue(result.containsKey(metadataFilename)); + } + public void testMarkMergedSegmentPendingDownload() { String localSegmentName1 = "_1.si"; String remoteSegmentName1 = "_1.si__uuid"; diff --git a/server/src/test/java/org/opensearch/index/store/RemoteSegmentStoreDirectoryWithPinnedTimestampTests.java b/server/src/test/java/org/opensearch/index/store/RemoteSegmentStoreDirectoryWithPinnedTimestampTests.java index 48a71c55bb8ce..63c3fd332bf28 100644 --- a/server/src/test/java/org/opensearch/index/store/RemoteSegmentStoreDirectoryWithPinnedTimestampTests.java +++ b/server/src/test/java/org/opensearch/index/store/RemoteSegmentStoreDirectoryWithPinnedTimestampTests.java @@ -35,9 +35,10 @@ import static org.opensearch.indices.RemoteStoreSettings.CLUSTER_REMOTE_STORE_PINNED_TIMESTAMP_ENABLED; import static org.hamcrest.CoreMatchers.is; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.argThat; +import static org.mockito.Mockito.any; import static org.mockito.Mockito.anyInt; +import static org.mockito.Mockito.anyString; +import static org.mockito.Mockito.argThat; import static org.mockito.Mockito.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; @@ -178,8 +179,8 @@ public void testDeleteStaleCommitsNoPinnedTimestampMdFilesLatest() throws Except remoteSegmentStoreDirectory.deleteStaleSegmentsAsync(2); assertBusy(() -> assertThat(remoteSegmentStoreDirectory.canDeleteStaleCommits.get(), is(true))); - verify(remoteDataDirectory, times(0)).deleteFile(any()); - verify(remoteMetadataDirectory, times(0)).deleteFile(any()); + verify(remoteDataDirectory, times(0)).deleteFile(anyString()); + verify(remoteMetadataDirectory, times(0)).deleteFile(anyString()); } public void testDeleteStaleCommitsPinnedTimestampMdFile() throws Exception { diff --git a/server/src/test/java/org/opensearch/index/store/checksum/ChecksumHandlerTests.java b/server/src/test/java/org/opensearch/index/store/checksum/ChecksumHandlerTests.java new file mode 100644 index 0000000000000..56d8682eccc50 --- /dev/null +++ b/server/src/test/java/org/opensearch/index/store/checksum/ChecksumHandlerTests.java @@ -0,0 +1,128 @@ +/* + * 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.store.checksum; + +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.store.ByteBuffersDirectory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexOutput; +import org.opensearch.index.store.FormatChecksumStrategy; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.zip.CRC32; + +/** + * Unit tests for {@link FormatChecksumStrategy} implementations: + * {@link LuceneChecksumHandler} and {@link GenericCRC32ChecksumHandler}. + */ +public class ChecksumHandlerTests extends OpenSearchTestCase { + + // ═══════════════════════════════════════════════════════════════ + // GenericCRC32ChecksumHandler Tests + // ═══════════════════════════════════════════════════════════════ + + public void testGenericCRC32ChecksumHandler_EmptyFile() throws IOException { + GenericCRC32ChecksumHandler handler = new GenericCRC32ChecksumHandler(); + ByteBuffersDirectory dir = new ByteBuffersDirectory(); + + try (IndexOutput out = dir.createOutput("empty.dat", IOContext.DEFAULT)) { + // write nothing + } + + long checksum = handler.computeChecksum(dir, "empty.dat"); + CRC32 crc32 = new CRC32(); + assertEquals("CRC32 of empty file should match", crc32.getValue(), checksum); + } + + public void testGenericCRC32ChecksumHandler_LargeFile() throws IOException { + GenericCRC32ChecksumHandler handler = new GenericCRC32ChecksumHandler(); + ByteBuffersDirectory dir = new ByteBuffersDirectory(); + + // Write a file larger than the 8192 buffer size + byte[] data = new byte[20000]; + for (int i = 0; i < data.length; i++) { + data[i] = (byte) (i % 256); + } + + try (IndexOutput out = dir.createOutput("large.dat", IOContext.DEFAULT)) { + out.writeBytes(data, data.length); + } + + long checksum = handler.computeChecksum(dir, "large.dat"); + + CRC32 crc32 = new CRC32(); + crc32.update(data); + assertEquals("CRC32 of large file should match", crc32.getValue(), checksum); + } + + // ═══════════════════════════════════════════════════════════════ + // LuceneChecksumHandler Tests + // ═══════════════════════════════════════════════════════════════ + + public void testLuceneChecksumHandler_ComputeChecksum() throws IOException { + LuceneChecksumHandler handler = new LuceneChecksumHandler(); + ByteBuffersDirectory dir = new ByteBuffersDirectory(); + + try (IndexOutput out = dir.createOutput("lucene_test.si", IOContext.DEFAULT)) { + CodecUtil.writeHeader(out, "TestCodec", 1); + out.writeString("some lucene data"); + CodecUtil.writeFooter(out); + } + + long checksum = handler.computeChecksum(dir, "lucene_test.si"); + assertTrue("Lucene checksum should be non-zero", checksum != 0); + } + + // ═══════════════════════════════════════════════════════════════ + // FormatChecksumStrategy default method Tests + // ═══════════════════════════════════════════════════════════════ + + public void testComputeChecksum_UploadChecksumString() throws IOException { + GenericCRC32ChecksumHandler handler = new GenericCRC32ChecksumHandler(); + ByteBuffersDirectory dir = new ByteBuffersDirectory(); + byte[] data = "test".getBytes(StandardCharsets.UTF_8); + + try (IndexOutput out = dir.createOutput("default_upload.dat", IOContext.DEFAULT)) { + out.writeBytes(data, data.length); + } + + long checksum = handler.computeChecksum(dir, "default_upload.dat"); + String uploadChecksum = Long.toString(checksum); + assertNotNull(uploadChecksum); + Long.parseLong(uploadChecksum); // should not throw + } + + public void testComputeChecksum_Idempotent() throws IOException { + GenericCRC32ChecksumHandler handler = new GenericCRC32ChecksumHandler(); + ByteBuffersDirectory dir = new ByteBuffersDirectory(); + byte[] data = "idempotent test".getBytes(StandardCharsets.UTF_8); + + try (IndexOutput out = dir.createOutput("idem.dat", IOContext.DEFAULT)) { + out.writeBytes(data, data.length); + } + + long checksum1 = handler.computeChecksum(dir, "idem.dat"); + long checksum2 = handler.computeChecksum(dir, "idem.dat"); + assertEquals("Checksum should be the same on repeated calls", checksum1, checksum2); + } + + public void testRegisterChecksum_DefaultNoOp() { + GenericCRC32ChecksumHandler handler = new GenericCRC32ChecksumHandler(); + // Default registerChecksum is a no-op; should not throw + handler.registerChecksum("file.dat", 12345L, 1L); + } + + public void testClearChecksums_DefaultNoOp() { + GenericCRC32ChecksumHandler handler = new GenericCRC32ChecksumHandler(); + // Default clearChecksums is a no-op; should not throw + handler.clearChecksums(); + } +} diff --git a/server/src/test/java/org/opensearch/index/store/remote/DataFormatAwareRemoteDirectoryTests.java b/server/src/test/java/org/opensearch/index/store/remote/DataFormatAwareRemoteDirectoryTests.java new file mode 100644 index 0000000000000..62a571aab9a41 --- /dev/null +++ b/server/src/test/java/org/opensearch/index/store/remote/DataFormatAwareRemoteDirectoryTests.java @@ -0,0 +1,1271 @@ +/* + * 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.store.remote; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.index.CorruptIndexException; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.IndexOutput; +import org.opensearch.Version; +import org.opensearch.cluster.metadata.IndexMetadata; +import org.opensearch.common.blobstore.AsyncMultiStreamBlobContainer; +import org.opensearch.common.blobstore.BlobContainer; +import org.opensearch.common.blobstore.BlobMetadata; +import org.opensearch.common.blobstore.BlobPath; +import org.opensearch.common.blobstore.BlobStore; +import org.opensearch.common.blobstore.exception.CorruptFileException; +import org.opensearch.common.blobstore.stream.write.WriteContext; +import org.opensearch.common.blobstore.support.PlainBlobMetadata; +import org.opensearch.common.blobstore.transfer.stream.OffsetRangeInputStream; +import org.opensearch.common.settings.Settings; +import org.opensearch.core.action.ActionListener; +import org.opensearch.index.IndexSettings; +import org.opensearch.index.engine.dataformat.DataFormatDescriptor; +import org.opensearch.index.engine.dataformat.DataFormatRegistry; +import org.opensearch.index.store.DataFormatAwareStoreDirectory; +import org.opensearch.index.store.FileMetadata; +import org.opensearch.index.store.RemoteIndexOutput; +import org.opensearch.index.store.RemoteSegmentStoreDirectory.UploadedSegmentMetadata; +import org.opensearch.index.store.checksum.GenericCRC32ChecksumHandler; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.nio.file.NoSuchFileException; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.UnaryOperator; + +import org.mockito.Mockito; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Unit tests for {@link DataFormatAwareRemoteDirectory}. + */ +public class DataFormatAwareRemoteDirectoryTests extends OpenSearchTestCase { + + private static final Logger logger = LogManager.getLogger(DataFormatAwareRemoteDirectoryTests.class); + + private BlobStore mockBlobStore; + private BlobContainer baseBlobContainer; + private BlobContainer parquetBlobContainer; + private BlobPath baseBlobPath; + private DataFormatAwareRemoteDirectory directory; + + @Override + public void setUp() throws Exception { + super.setUp(); + mockBlobStore = mock(BlobStore.class); + baseBlobContainer = mock(BlobContainer.class); + parquetBlobContainer = mock(BlobContainer.class); + baseBlobPath = new BlobPath().add("segments").add("data"); + + when(mockBlobStore.blobContainer(baseBlobPath)).thenReturn(baseBlobContainer); + when(mockBlobStore.blobContainer(baseBlobPath.add("parquet"))).thenReturn(parquetBlobContainer); + + // Identity rate limiters (no-op) + UnaryOperator uploadRateLimiter = UnaryOperator.identity(); + UnaryOperator lowPriorityUploadRateLimiter = UnaryOperator.identity(); + UnaryOperator downloadRateLimiter = UnaryOperator.identity(); + UnaryOperator lowPriorityDownloadRateLimiter = UnaryOperator.identity(); + + // Mock DataFormatRegistry to register "parquet" format + DataFormatRegistry mockRegistry = mock(DataFormatRegistry.class); + Settings indexSettingsBuilder = Settings.builder() + .put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT) + .put(IndexMetadata.SETTING_INDEX_UUID, "test-uuid") + .build(); + IndexMetadata metadata = IndexMetadata.builder("test-index") + .settings(indexSettingsBuilder) + .numberOfShards(1) + .numberOfReplicas(0) + .build(); + IndexSettings indexSettings = new IndexSettings(metadata, Settings.EMPTY); + when(mockRegistry.getFormatDescriptors(any(IndexSettings.class))).thenReturn( + Map.of("parquet", new DataFormatDescriptor("parquet", new GenericCRC32ChecksumHandler())) + ); + + directory = new DataFormatAwareRemoteDirectory( + mockBlobStore, + baseBlobPath, + uploadRateLimiter, + lowPriorityUploadRateLimiter, + downloadRateLimiter, + lowPriorityDownloadRateLimiter, + new HashMap<>(), + logger, + mockRegistry, + indexSettings + ); + } + + // ═══════════════════════════════════════════════════════════════ + // Format Routing Tests (getBlobContainerForFormat) + // ═══════════════════════════════════════════════════════════════ + + public void testGetBlobContainerForFormat_Lucene() { + BlobContainer container = directory.getBlobContainerForFormat("lucene"); + assertSame("lucene should route to base container", baseBlobContainer, container); + } + + public void testGetBlobContainerForFormat_LUCENE_UpperCase() { + BlobContainer container = directory.getBlobContainerForFormat("LUCENE"); + assertSame("LUCENE should route to base container", baseBlobContainer, container); + } + + public void testGetBlobContainerForFormat_Metadata() { + BlobContainer container = directory.getBlobContainerForFormat("metadata"); + assertSame("metadata should route to base container", baseBlobContainer, container); + } + + public void testGetBlobContainerForFormat_Null() { + BlobContainer container = directory.getBlobContainerForFormat(null); + assertSame("null format should route to base container", baseBlobContainer, container); + } + + public void testGetBlobContainerForFormat_Empty() { + BlobContainer container = directory.getBlobContainerForFormat(""); + assertSame("empty format should route to base container", baseBlobContainer, container); + } + + public void testGetBlobContainerForFormat_Parquet() { + BlobContainer container = directory.getBlobContainerForFormat("parquet"); + assertSame("parquet should route to parquet container", parquetBlobContainer, container); + } + + public void testGetBlobContainerForFormat_UnregisteredFormat_CreatesLazily() { + // FormatBlobRouter lazily creates containers for unknown formats via computeIfAbsent + // With a mock BlobStore, the container may be null, but no exception is thrown + assertNoException(() -> directory.getBlobContainerForFormat("arrow")); + } + + private static void assertNoException(Runnable r) { + r.run(); + } + + // ═══════════════════════════════════════════════════════════════ + // List Tests + // ═══════════════════════════════════════════════════════════════ + + public void testListAll_AggregatesAllContainers() throws IOException { + // Base container has lucene files + Map baseBlobs = new HashMap<>(); + baseBlobs.put("_0.cfs__UUID1", new PlainBlobMetadata("_0.cfs__UUID1", 100)); + baseBlobs.put("_0.si__UUID2", new PlainBlobMetadata("_0.si__UUID2", 50)); + when(baseBlobContainer.listBlobs()).thenReturn(baseBlobs); + + Map parquetBlobs = new HashMap<>(); + parquetBlobs.put("_0.parquet__UUID3", new PlainBlobMetadata("_0.parquet__UUID3", 200)); + when(parquetBlobContainer.listBlobs()).thenReturn(parquetBlobs); + + String[] allFiles = directory.listAll(); + + assertEquals(3, allFiles.length); + // Should be sorted + assertEquals("_0.cfs__UUID1", allFiles[0]); + assertEquals("_0.parquet__UUID3", allFiles[1]); + assertEquals("_0.si__UUID2", allFiles[2]); + } + + public void testListAll_EmptyContainers() throws IOException { + when(baseBlobContainer.listBlobs()).thenReturn(Collections.emptyMap()); + + String[] allFiles = directory.listAll(); + assertEquals(0, allFiles.length); + } + + // ═══════════════════════════════════════════════════════════════ + // Delete Tests + // ═══════════════════════════════════════════════════════════════ + + public void testDeleteFile_LuceneFile() throws IOException { + directory.deleteFile("_0.cfs"); + + verify(baseBlobContainer).deleteBlobsIgnoringIfNotExists(Collections.singletonList("_0.cfs")); + verify(parquetBlobContainer, never()).deleteBlobsIgnoringIfNotExists(any()); + } + + public void testDeleteFile_ParquetFile_WithFormatSuffix() throws IOException { + // Register format in cache since DFARD receives plain blob keys + directory.getFormatBlobRouter().orElseThrow().registerBlobFormat("_0.parquet", "parquet"); + directory.deleteFile("_0.parquet"); + + verify(parquetBlobContainer).deleteBlobsIgnoringIfNotExists(Collections.singletonList("_0.parquet")); + } + + public void testDeleteFile_WithUploadedSegmentMetadata_Parquet() throws IOException { + + // Use fromString() since constructor is package-private + // Format: originalFilename::uploadedFilename::checksum::length::writtenByMajor + UploadedSegmentMetadata metadata = UploadedSegmentMetadata.fromString( + "parquet/_0.parquet::_0.parquet__UUID1::checksum123::200::10" + ); + + directory.getFormatBlobRouter().orElseThrow().registerBlobFormat("_0.parquet__UUID1", "parquet"); + directory.deleteFile(metadata.getUploadedFilename()); + + verify(parquetBlobContainer).deleteBlobsIgnoringIfNotExists(Collections.singletonList("_0.parquet__UUID1")); + } + + public void testDeleteFiles_BatchDelete_DeletesFromAllContainers() throws IOException { + + List names = List.of("_0.cfs__UUID1", "_0.parquet__UUID2"); + directory.deleteFiles(names); + + // baseBlobContainer is called from super.deleteFiles + lucene format container (same instance) + verify(baseBlobContainer, times(2)).deleteBlobsIgnoringIfNotExists(names); + verify(parquetBlobContainer).deleteBlobsIgnoringIfNotExists(names); + } + + public void testDeleteFiles_EmptyList_NoOp() throws IOException { + directory.deleteFiles(Collections.emptyList()); + + verify(baseBlobContainer, never()).deleteBlobsIgnoringIfNotExists(any()); + } + + public void testDeleteFiles_NullList_NoOp() throws IOException { + directory.deleteFiles(null); + + verify(baseBlobContainer, never()).deleteBlobsIgnoringIfNotExists(any()); + } + + // ═══════════════════════════════════════════════════════════════ + // OpenInput Tests + // ═══════════════════════════════════════════════════════════════ + + public void testOpenInput_WithUploadedSegmentMetadata_Lucene() throws IOException { + UploadedSegmentMetadata metadata = UploadedSegmentMetadata.fromString("_0.cfs::_0.cfs__UUID1::checksum123::100::10"); + + byte[] content = new byte[100]; + when(baseBlobContainer.readBlob("_0.cfs__UUID1")).thenReturn(new ByteArrayInputStream(content)); + + IndexInput input = directory.openInput(metadata.getUploadedFilename(), 100, IOContext.DEFAULT); + assertNotNull(input); + assertEquals(100, input.length()); + input.close(); + + verify(baseBlobContainer).readBlob("_0.cfs__UUID1"); + verify(parquetBlobContainer, never()).readBlob(anyString()); + } + + public void testOpenInput_WithUploadedSegmentMetadata_Parquet() throws IOException { + + UploadedSegmentMetadata metadata = UploadedSegmentMetadata.fromString( + "parquet/_0.parquet::_0.parquet__UUID1::checksum456::200::10" + ); + + byte[] content = new byte[200]; + when(parquetBlobContainer.readBlob("_0.parquet__UUID1")).thenReturn(new ByteArrayInputStream(content)); + + directory.getFormatBlobRouter().orElseThrow().registerBlobFormat("_0.parquet__UUID1", "parquet"); + IndexInput input = directory.openInput(metadata.getUploadedFilename(), 200, IOContext.DEFAULT); + assertNotNull(input); + assertEquals(200, input.length()); + input.close(); + + verify(parquetBlobContainer).readBlob("_0.parquet__UUID1"); + verify(baseBlobContainer, never()).readBlob(anyString()); + } + + public void testOpenInput_ClosesStream_OnFailure() throws IOException { + InputStream mockStream = mock(InputStream.class); + when(baseBlobContainer.readBlob("_0.cfs__UUID1")).thenReturn(mockStream); + when(mockStream.read(any(), anyInt(), anyInt())).thenThrow(new IOException("read error")); + + UploadedSegmentMetadata metadata = UploadedSegmentMetadata.fromString("_0.cfs::_0.cfs__UUID1::checksum123::100::10"); + + // The openInput should succeed (it just wraps the stream), but we verify the pattern + IndexInput input = directory.openInput(metadata.getUploadedFilename(), 100, IOContext.DEFAULT); + assertNotNull(input); + input.close(); + } + + // ═══════════════════════════════════════════════════════════════ + // FileLength Tests + // ═══════════════════════════════════════════════════════════════ + + public void testFileLength_LuceneFile() throws IOException { + List blobList = List.of(new PlainBlobMetadata("_0.cfs", 1234)); + when(baseBlobContainer.listBlobsByPrefixInSortedOrder(eq("_0.cfs"), eq(1), any())).thenReturn(blobList); + + long length = directory.fileLength("_0.cfs"); + assertEquals(1234, length); + } + + public void testFileLength_ParquetFile() throws IOException { + + List blobList = List.of(new PlainBlobMetadata("_0.parquet", 5678)); + when(parquetBlobContainer.listBlobsByPrefixInSortedOrder(eq("_0.parquet"), eq(1), any())).thenReturn(blobList); + + // Register format so resolveFormat routes to parquet container + directory.getFormatBlobRouter().orElseThrow().registerBlobFormat("_0.parquet", "parquet"); + long length = directory.fileLength("_0.parquet"); + assertEquals(5678, length); + } + + public void testFileLength_FileNotFound() throws IOException { + when(baseBlobContainer.listBlobsByPrefixInSortedOrder(eq("nonexistent"), eq(1), any())).thenReturn(Collections.emptyList()); + + expectThrows(NoSuchFileException.class, () -> directory.fileLength("nonexistent")); + } + + // ═══════════════════════════════════════════════════════════════ + // Lifecycle Tests + // ═══════════════════════════════════════════════════════════════ + + public void testDelete_DeletesAllContainers() throws IOException { + + directory.delete(); + + // baseBlobContainer is used for both the "lucene" format and the inherited base container + verify(parquetBlobContainer).delete(); + verify(baseBlobContainer, times(2)).delete(); + } + + public void testClose_ClearsFormatContainers() throws IOException { + // Verify parquet container exists before close + assertNotNull(directory.getBlobContainerForFormat("parquet")); + + directory.close(); + + // After close, the directory is closed but FormatBlobRouter still lazily creates containers + // The important thing is that close() doesn't throw + } + + // ═══════════════════════════════════════════════════════════════ + // Edge Case Tests + // ═══════════════════════════════════════════════════════════════ + + public void testConstructor_NullDataFormatRegistry() { + // Should not throw with null DataFormatRegistry and IndexSettings + DataFormatAwareRemoteDirectory dir = new DataFormatAwareRemoteDirectory( + mockBlobStore, + baseBlobPath, + UnaryOperator.identity(), + UnaryOperator.identity(), + UnaryOperator.identity(), + UnaryOperator.identity(), + new HashMap<>(), + logger, + null, + null + ); + assertNotNull(dir); + } + + public void testToString() { + + String str = directory.toString(); + assertTrue(str.contains("DataFormatAwareRemoteDirectory")); + assertTrue(str.contains("parquet")); + } + + // ═══════════════════════════════════════════════════════════════ + // Sync CopyFrom Tests + // ═══════════════════════════════════════════════════════════════ + + public void testSyncCopyFrom_RoutesToCorrectContainer() throws IOException { + // We can't easily test the full copyFrom without a real Directory, + // but we can verify that the format routing logic works by testing + // the getBlobContainerForFormat that copyFrom uses internally. + + // For lucene files + BlobContainer luceneContainer = directory.getBlobContainerForFormat("lucene"); + assertSame(baseBlobContainer, luceneContainer); + + // For parquet files + BlobContainer parquetContainer = directory.getBlobContainerForFormat("parquet"); + assertSame(parquetBlobContainer, parquetContainer); + } + + // ═══════════════════════════════════════════════════════════════ + // openInput(String, long, IOContext) Tests - Format-aware routing + // ═══════════════════════════════════════════════════════════════ + + public void testOpenInput_StringBased_LuceneFile() throws IOException { + byte[] content = new byte[100]; + when(baseBlobContainer.readBlob("_0.cfs")).thenReturn(new ByteArrayInputStream(content)); + + IndexInput input = directory.openInput("_0.cfs", 100, IOContext.DEFAULT); + assertNotNull(input); + assertEquals(100, input.length()); + input.close(); + + verify(baseBlobContainer).readBlob("_0.cfs"); + } + + public void testOpenInput_StringBased_ParquetFile_WithFormatSuffix() throws IOException { + + byte[] content = new byte[200]; + when(parquetBlobContainer.readBlob("_0.parquet")).thenReturn(new ByteArrayInputStream(content)); + + // Register format in cache since DFARD receives plain blob keys + directory.getFormatBlobRouter().orElseThrow().registerBlobFormat("_0.parquet", "parquet"); + IndexInput input = directory.openInput("_0.parquet", 200, IOContext.DEFAULT); + assertNotNull(input); + assertEquals(200, input.length()); + input.close(); + + verify(parquetBlobContainer).readBlob("_0.parquet"); + verify(baseBlobContainer, never()).readBlob(anyString()); + } + + public void testOpenInput_StringBased_ClosesStreamOnException() throws IOException { + when(baseBlobContainer.readBlob("_0.cfs")).thenThrow(new IOException("read error")); + + expectThrows(IOException.class, () -> directory.openInput("_0.cfs", 100, IOContext.DEFAULT)); + } + + // ═══════════════════════════════════════════════════════════════ + // createOutput Tests - Format-aware routing + // ═══════════════════════════════════════════════════════════════ + + public void testCreateOutput_LuceneFormat() throws IOException { + RemoteIndexOutput output = directory.createOutput("test_file", "lucene", IOContext.DEFAULT); + assertNotNull(output); + } + + public void testCreateOutput_ParquetFormat() throws IOException { + + RemoteIndexOutput output = directory.createOutput("test_file.parquet", "parquet", IOContext.DEFAULT); + assertNotNull(output); + } + + // ═══════════════════════════════════════════════════════════════ + // fileLength(FileMetadata) Tests + // ═══════════════════════════════════════════════════════════════ + + public void testFileLength_FileMetadata_Lucene() throws IOException { + FileMetadata fm = new FileMetadata("lucene", "_0.cfs"); + List blobList = List.of(new PlainBlobMetadata("_0.cfs", 1234)); + when(baseBlobContainer.listBlobsByPrefixInSortedOrder(eq("_0.cfs"), eq(1), any())).thenReturn(blobList); + + long length = directory.fileLength(fm.file()); + assertEquals(1234, length); + } + + public void testFileLength_FileMetadata_Parquet() throws IOException { + FileMetadata fm = new FileMetadata("parquet", "_0.parquet"); + List blobList = List.of(new PlainBlobMetadata("_0.parquet", 5678)); + when(parquetBlobContainer.listBlobsByPrefixInSortedOrder(eq("_0.parquet"), eq(1), any())).thenReturn(blobList); + + directory.getFormatBlobRouter().orElseThrow().registerBlobFormat("_0.parquet", "parquet"); + long length = directory.fileLength(fm.file()); + assertEquals(5678, length); + } + + public void testFileLength_FileMetadata_NotFound() throws IOException { + FileMetadata fm = new FileMetadata("lucene", "nonexistent"); + when(baseBlobContainer.listBlobsByPrefixInSortedOrder(eq("nonexistent"), eq(1), any())).thenReturn(Collections.emptyList()); + + expectThrows(NoSuchFileException.class, () -> directory.fileLength(fm.file())); + } + + // ═══════════════════════════════════════════════════════════════ + // openInput(FileMetadata, long, IOContext) Tests + // ═══════════════════════════════════════════════════════════════ + + public void testOpenInput_FileMetadata_Lucene() throws IOException { + FileMetadata fm = new FileMetadata("lucene", "_0.cfs"); + byte[] content = new byte[100]; + when(baseBlobContainer.readBlob("_0.cfs")).thenReturn(new ByteArrayInputStream(content)); + + IndexInput input = directory.openInput(fm.file(), 100, IOContext.DEFAULT); + assertNotNull(input); + assertEquals(100, input.length()); + input.close(); + + verify(baseBlobContainer).readBlob("_0.cfs"); + } + + public void testOpenInput_FileMetadata_Parquet() throws IOException { + FileMetadata fm = new FileMetadata("parquet", "_0.parquet"); + byte[] content = new byte[200]; + when(parquetBlobContainer.readBlob("_0.parquet")).thenReturn(new ByteArrayInputStream(content)); + + directory.getFormatBlobRouter().orElseThrow().registerBlobFormat("_0.parquet", "parquet"); + IndexInput input = directory.openInput(fm.file(), 200, IOContext.DEFAULT); + assertNotNull(input); + assertEquals(200, input.length()); + input.close(); + + verify(parquetBlobContainer).readBlob("_0.parquet"); + verify(baseBlobContainer, never()).readBlob(anyString()); + } + + public void testOpenInput_FileMetadata_ClosesStreamOnException() throws IOException { + FileMetadata fm = new FileMetadata("lucene", "_0.cfs"); + InputStream mockStream = mock(InputStream.class); + when(baseBlobContainer.readBlob("_0.cfs")).thenReturn(mockStream); + when(mockStream.read(any(), anyInt(), anyInt())).thenThrow(new IOException("read error")); + + // openInput wraps the stream, reading from it will fail + IndexInput input = directory.openInput(fm.file(), 100, IOContext.DEFAULT); + assertNotNull(input); + input.close(); + } + + // ═══════════════════════════════════════════════════════════════ + // deleteFile(UploadedSegmentMetadata) - Lucene metadata + // ═══════════════════════════════════════════════════════════════ + + public void testDeleteFile_WithUploadedSegmentMetadata_Lucene() throws IOException { + UploadedSegmentMetadata metadata = UploadedSegmentMetadata.fromString("_0.cfs::_0.cfs__UUID1::checksum123::100::10"); + + directory.deleteFile(metadata.getUploadedFilename()); + + verify(baseBlobContainer).deleteBlobsIgnoringIfNotExists(Collections.singletonList("_0.cfs__UUID1")); + verify(parquetBlobContainer, never()).deleteBlobsIgnoringIfNotExists(any()); + } + + // ═══════════════════════════════════════════════════════════════ + // Async copyFrom Tests (8-arg version, returns boolean) + // ═══════════════════════════════════════════════════════════════ + + public void testAsyncCopyFrom_NonAsyncContainer_ReturnsFalse() throws IOException { + // baseBlobContainer is NOT AsyncMultiStreamBlobContainer, so should return false + org.opensearch.core.action.ActionListener listener = mock(org.opensearch.core.action.ActionListener.class); + org.apache.lucene.store.Directory mockFrom = mock(org.apache.lucene.store.Directory.class); + + boolean result = directory.copyFrom( + mockFrom, + "_0.cfs", // src (lucene format, no :::) + "_0.cfs__UUID", // remoteFileName + IOContext.DEFAULT, + () -> {}, + listener, + false, + null + ); + + assertFalse("Should return false when container is not AsyncMultiStreamBlobContainer", result); + } + + public void testAsyncCopyFrom_ExceptionHandling() throws IOException { + org.opensearch.core.action.ActionListener listener = mock(org.opensearch.core.action.ActionListener.class); + org.apache.lucene.store.Directory mockFrom = mock(org.apache.lucene.store.Directory.class); + // Make openInput throw an exception + when(mockFrom.openInput(anyString(), any())).thenThrow(new IOException("open failed")); + + // Even with exception, it should not propagate but call listener.onFailure + boolean result = directory.copyFrom(mockFrom, "_0.cfs", "_0.cfs__UUID", IOContext.DEFAULT, () -> {}, listener, false, null); + + // Returns false because baseBlobContainer is not AsyncMultiStreamBlobContainer + assertFalse(result); + } + + // ═══════════════════════════════════════════════════════════════ + // FileMetadata-based copyFrom (non-async) - returns false when not async + // ═══════════════════════════════════════════════════════════════ + + public void testCopyFrom_FileMetadata_NonAsync_ReturnsFalse() throws IOException { + org.opensearch.core.action.ActionListener listener = mock(org.opensearch.core.action.ActionListener.class); + org.apache.lucene.store.Directory mockFrom = mock(org.apache.lucene.store.Directory.class); + + boolean result = directory.copyFrom( + mockFrom, + "_0.cfs:::lucene", + "_0.cfs__UUID", + IOContext.DEFAULT, + () -> {}, + listener, + false, + null + ); + + assertFalse("Should return false when base container is not AsyncMultiStreamBlobContainer", result); + } + + // ═══════════════════════════════════════════════════════════════ + // openInput(UploadedSegmentMetadata) - Exception closes stream + // ═══════════════════════════════════════════════════════════════ + + public void testOpenInput_UploadedSegmentMetadata_ExceptionClosesStream() throws IOException { + UploadedSegmentMetadata metadata = UploadedSegmentMetadata.fromString("_0.cfs::_0.cfs__UUID1::checksum123::100::10"); + when(baseBlobContainer.readBlob("_0.cfs__UUID1")).thenThrow(new IOException("blob read failed")); + + expectThrows(IOException.class, () -> directory.openInput(metadata.getUploadedFilename(), 100, IOContext.DEFAULT)); + } + + // ═══════════════════════════════════════════════════════════════ + // Metadata file routing + // ═══════════════════════════════════════════════════════════════ + + public void testDeleteFile_MetadataFile() throws IOException { + directory.deleteFile("metadata__1__2__3"); + + // "metadata" format routes to base container + verify(baseBlobContainer).deleteBlobsIgnoringIfNotExists(Collections.singletonList("metadata__1__2__3")); + } + + // ═══════════════════════════════════════════════════════════════ + // Sync copyFrom(Directory, String, String, IOContext) Tests + // ═══════════════════════════════════════════════════════════════ + + public void testSyncCopyFrom_LuceneFile_CopiesToBaseContainer() throws IOException { + Directory mockFrom = mock(Directory.class); + IndexInput mockInput = mock(IndexInput.class); + when(mockInput.length()).thenReturn(10L); + when(mockFrom.openInput(eq("_0.cfs"), any(IOContext.class))).thenReturn(mockInput); + + // The copyFrom creates a RemoteIndexOutput that writes to the base container + directory.copyFrom(mockFrom, "_0.cfs", "_0.cfs__UUID", IOContext.DEFAULT); + + verify(mockFrom).openInput(eq("_0.cfs"), any(IOContext.class)); + } + + public void testSyncCopyFrom_ParquetFile_CopiesToFormatContainer() throws IOException { + Directory mockFrom = mock(Directory.class); + IndexInput mockInput = mock(IndexInput.class); + when(mockInput.length()).thenReturn(20L); + when(mockFrom.openInput(eq("_0.pqt:::parquet"), any(IOContext.class))).thenReturn(mockInput); + + directory.copyFrom(mockFrom, "_0.pqt:::parquet", "_0.pqt__UUID", IOContext.DEFAULT); + + verify(mockFrom).openInput(eq("_0.pqt:::parquet"), any(IOContext.class)); + } + + // ═══════════════════════════════════════════════════════════════ + // Async copyFrom with AsyncMultiStreamBlobContainer Tests + // ═══════════════════════════════════════════════════════════════ + + public void testAsyncCopyFrom_WithAsyncContainer_ReturnsTrue() throws Exception { + // Create an async blob container + AsyncMultiStreamBlobContainer asyncContainer = mock(AsyncMultiStreamBlobContainer.class); + when(asyncContainer.remoteIntegrityCheckSupported()).thenReturn(false); + when(asyncContainer.path()).thenReturn(baseBlobPath); + + // Wire up blobStore to return async container for base path + BlobStore asyncBlobStore = mock(BlobStore.class); + when(asyncBlobStore.blobContainer(baseBlobPath)).thenReturn(asyncContainer); + + DataFormatAwareRemoteDirectory asyncDir = new DataFormatAwareRemoteDirectory( + asyncBlobStore, + baseBlobPath, + UnaryOperator.identity(), + UnaryOperator.identity(), + UnaryOperator.identity(), + UnaryOperator.identity(), + new HashMap<>(), + logger, + null, + null + ); + + // Set up async upload to call onResponse + Mockito.doAnswer(invocation -> { + ActionListener completionListener = invocation.getArgument(1); + completionListener.onResponse(null); + return null; + }).when(asyncContainer).asyncBlobUpload(any(WriteContext.class), any()); + + // Create a real directory with a file that has a valid codec footer + Directory storeDirectory = newDirectory(); + String filename = "_100.si"; + IndexOutput indexOutput = storeDirectory.createOutput(filename, IOContext.DEFAULT); + indexOutput.writeString("Hello World!"); + CodecUtil.writeFooter(indexOutput); + indexOutput.close(); + storeDirectory.sync(List.of(filename)); + + CountDownLatch latch = new CountDownLatch(1); + AtomicReference postUploadInvoked = new AtomicReference<>(false); + + boolean result = asyncDir.copyFrom( + storeDirectory, + filename, + filename, + IOContext.DEFAULT, + () -> postUploadInvoked.set(true), + new ActionListener<>() { + @Override + public void onResponse(Void unused) { + latch.countDown(); + } + + @Override + public void onFailure(Exception e) { + fail("Should not fail: " + e.getMessage()); + } + }, + false, + null + ); + + assertTrue("Should return true when container is AsyncMultiStreamBlobContainer", result); + assertTrue(latch.await(10, TimeUnit.SECONDS)); + assertTrue(postUploadInvoked.get()); + storeDirectory.close(); + } + + public void testAsyncCopyFrom_ExceptionDuringUpload_CallsListenerOnFailure() throws Exception { + AsyncMultiStreamBlobContainer asyncContainer = mock(AsyncMultiStreamBlobContainer.class); + when(asyncContainer.remoteIntegrityCheckSupported()).thenReturn(false); + when(asyncContainer.path()).thenReturn(baseBlobPath); + + BlobStore asyncBlobStore = mock(BlobStore.class); + when(asyncBlobStore.blobContainer(baseBlobPath)).thenReturn(asyncContainer); + + DataFormatAwareRemoteDirectory asyncDir = new DataFormatAwareRemoteDirectory( + asyncBlobStore, + baseBlobPath, + UnaryOperator.identity(), + UnaryOperator.identity(), + UnaryOperator.identity(), + UnaryOperator.identity(), + new HashMap<>(), + logger, + null, + null + ); + + // File does not exist - openInput will throw + Directory storeDirectory = newDirectory(); + CountDownLatch latch = new CountDownLatch(1); + AtomicReference failureRef = new AtomicReference<>(); + + boolean result = asyncDir.copyFrom( + storeDirectory, + "_nonexistent.si", + "_nonexistent.si__UUID", + IOContext.DEFAULT, + () -> {}, + new ActionListener<>() { + @Override + public void onResponse(Void unused) { + fail("Should have failed"); + } + + @Override + public void onFailure(Exception e) { + failureRef.set(e); + latch.countDown(); + } + }, + false, + null + ); + + assertTrue("Should return true (handled)", result); + assertTrue(latch.await(10, TimeUnit.SECONDS)); + assertNotNull(failureRef.get()); + storeDirectory.close(); + } + + // ═══════════════════════════════════════════════════════════════ + // openInput exception close-on-failure (stream opened but readBlob fails) + // ═══════════════════════════════════════════════════════════════ + + public void testOpenInput_FileMetadata_ExceptionClosesStream() throws IOException { + FileMetadata fm = new FileMetadata("lucene", "_0.cfs"); + when(baseBlobContainer.readBlob("_0.cfs")).thenThrow(new IOException("blob read failed")); + + expectThrows(IOException.class, () -> directory.openInput(fm.file(), 100, IOContext.DEFAULT)); + } + + public void testOpenInput_StringBased_StreamClosedWhenInputStreamReadFails() throws IOException { + InputStream mockStream = mock(InputStream.class); + when(baseBlobContainer.readBlob("_0.cfs")).thenReturn(mockStream); + when(mockStream.read(any(), anyInt(), anyInt())).thenThrow(new IOException("stream error")); + + // openInput wraps the stream successfully + IndexInput input = directory.openInput("_0.cfs", 100, IOContext.DEFAULT); + assertNotNull(input); + input.close(); + } + + // ═══════════════════════════════════════════════════════════════ + // fileLength name mismatch tests + // ═══════════════════════════════════════════════════════════════ + + public void testFileLength_NameMismatch_ThrowsNoSuchFile() throws IOException { + // Returns a blob but name doesn't match + List blobList = List.of(new PlainBlobMetadata("_0.cfs_DIFFERENT", 1234)); + when(baseBlobContainer.listBlobsByPrefixInSortedOrder(eq("_0.cfs"), eq(1), any())).thenReturn(blobList); + + expectThrows(NoSuchFileException.class, () -> directory.fileLength("_0.cfs")); + } + + public void testFileLength_FileMetadata_NameMismatch_ThrowsNoSuchFile() throws IOException { + FileMetadata fm = new FileMetadata("lucene", "_0.cfs"); + List blobList = List.of(new PlainBlobMetadata("_0.cfs_DIFFERENT", 1234)); + when(baseBlobContainer.listBlobsByPrefixInSortedOrder(eq("_0.cfs"), eq(1), any())).thenReturn(blobList); + + expectThrows(NoSuchFileException.class, () -> directory.fileLength(fm.file())); + } + + // ═══════════════════════════════════════════════════════════════ + // DownloadRateLimiterProvider merged segment path + // ═══════════════════════════════════════════════════════════════ + + public void testDownloadRateLimiter_MergedSegment_UsesLowPriorityRateLimiter() throws IOException { + // Create directory with pending merged segments + Map pendingMergedSegments = new HashMap<>(); + pendingMergedSegments.put("localFile", "_0.cfs__UUID_MERGED"); + + UnaryOperator normalRateLimiter = stream -> stream; + UnaryOperator lowPriorityRateLimiter = stream -> stream; + + DataFormatAwareRemoteDirectory dirWithMerged = new DataFormatAwareRemoteDirectory( + mockBlobStore, + baseBlobPath, + UnaryOperator.identity(), + UnaryOperator.identity(), + normalRateLimiter, + lowPriorityRateLimiter, + pendingMergedSegments, + logger, + null, + null + ); + + // When opening a merged segment, the low-priority rate limiter should be used + byte[] content = new byte[100]; + when(baseBlobContainer.readBlob("_0.cfs__UUID_MERGED")).thenReturn(new ByteArrayInputStream(content)); + + UploadedSegmentMetadata metadata = UploadedSegmentMetadata.fromString("_0.cfs::_0.cfs__UUID_MERGED::checksum123::100::10"); + + IndexInput input = dirWithMerged.openInput(metadata.getUploadedFilename(), 100, IOContext.DEFAULT); + assertNotNull(input); + assertEquals(100, input.length()); + input.close(); + } + + // ═══════════════════════════════════════════════════════════════ + // Completion listener tests (via async upload with errors) + // ═══════════════════════════════════════════════════════════════ + + public void testCompletionListener_PostUploadRunnerException() throws Exception { + AsyncMultiStreamBlobContainer asyncContainer = mock(AsyncMultiStreamBlobContainer.class); + when(asyncContainer.remoteIntegrityCheckSupported()).thenReturn(false); + when(asyncContainer.path()).thenReturn(baseBlobPath); + + BlobStore asyncBlobStore = mock(BlobStore.class); + when(asyncBlobStore.blobContainer(baseBlobPath)).thenReturn(asyncContainer); + + DataFormatAwareRemoteDirectory asyncDir = new DataFormatAwareRemoteDirectory( + asyncBlobStore, + baseBlobPath, + UnaryOperator.identity(), + UnaryOperator.identity(), + UnaryOperator.identity(), + UnaryOperator.identity(), + new HashMap<>(), + logger, + null, + null + ); + + Mockito.doAnswer(invocation -> { + ActionListener completionListener = invocation.getArgument(1); + completionListener.onResponse(null); + return null; + }).when(asyncContainer).asyncBlobUpload(any(WriteContext.class), any()); + + Directory storeDirectory = newDirectory(); + String filename = "_100.si"; + IndexOutput indexOutput = storeDirectory.createOutput(filename, IOContext.DEFAULT); + indexOutput.writeString("data"); + CodecUtil.writeFooter(indexOutput); + indexOutput.close(); + storeDirectory.sync(List.of(filename)); + + CountDownLatch latch = new CountDownLatch(1); + + // postUploadRunner throws exception → listener.onFailure + boolean result = asyncDir.copyFrom(storeDirectory, filename, filename + "__UUID", IOContext.DEFAULT, () -> { + throw new RuntimeException("postUpload error"); + }, new ActionListener<>() { + @Override + public void onResponse(Void unused) { + fail("Should not succeed"); + } + + @Override + public void onFailure(Exception e) { + latch.countDown(); + } + }, false, null); + + assertTrue(result); + assertTrue(latch.await(10, TimeUnit.SECONDS)); + storeDirectory.close(); + } + + public void testCompletionListener_CorruptIndexException() throws Exception { + AsyncMultiStreamBlobContainer asyncContainer = mock(AsyncMultiStreamBlobContainer.class); + when(asyncContainer.remoteIntegrityCheckSupported()).thenReturn(false); + when(asyncContainer.path()).thenReturn(baseBlobPath); + + BlobStore asyncBlobStore = mock(BlobStore.class); + when(asyncBlobStore.blobContainer(baseBlobPath)).thenReturn(asyncContainer); + + DataFormatAwareRemoteDirectory asyncDir = new DataFormatAwareRemoteDirectory( + asyncBlobStore, + baseBlobPath, + UnaryOperator.identity(), + UnaryOperator.identity(), + UnaryOperator.identity(), + UnaryOperator.identity(), + new HashMap<>(), + logger, + null, + null + ); + + // asyncBlobUpload calls onFailure with a wrapped CorruptIndexException + Mockito.doAnswer(invocation -> { + ActionListener completionListener = invocation.getArgument(1); + completionListener.onFailure(new RuntimeException(new CorruptIndexException("corrupted", "test"))); + return null; + }).when(asyncContainer).asyncBlobUpload(any(WriteContext.class), any()); + + Directory storeDirectory = newDirectory(); + String filename = "_100.si"; + IndexOutput indexOutput = storeDirectory.createOutput(filename, IOContext.DEFAULT); + indexOutput.writeString("data"); + CodecUtil.writeFooter(indexOutput); + indexOutput.close(); + storeDirectory.sync(List.of(filename)); + + CountDownLatch latch = new CountDownLatch(1); + AtomicReference failureRef = new AtomicReference<>(); + + boolean result = asyncDir.copyFrom( + storeDirectory, + filename, + filename + "__UUID", + IOContext.DEFAULT, + () -> {}, + new ActionListener<>() { + @Override + public void onResponse(Void unused) { + fail("Should not succeed"); + } + + @Override + public void onFailure(Exception e) { + failureRef.set(e); + latch.countDown(); + } + }, + false, + null + ); + + assertTrue(result); + assertTrue(latch.await(10, TimeUnit.SECONDS)); + assertTrue("Should be CorruptIndexException", failureRef.get() instanceof CorruptIndexException); + storeDirectory.close(); + } + + public void testCompletionListener_CorruptFileException() throws Exception { + AsyncMultiStreamBlobContainer asyncContainer = mock(AsyncMultiStreamBlobContainer.class); + when(asyncContainer.remoteIntegrityCheckSupported()).thenReturn(false); + when(asyncContainer.path()).thenReturn(baseBlobPath); + + BlobStore asyncBlobStore = mock(BlobStore.class); + when(asyncBlobStore.blobContainer(baseBlobPath)).thenReturn(asyncContainer); + + DataFormatAwareRemoteDirectory asyncDir = new DataFormatAwareRemoteDirectory( + asyncBlobStore, + baseBlobPath, + UnaryOperator.identity(), + UnaryOperator.identity(), + UnaryOperator.identity(), + UnaryOperator.identity(), + new HashMap<>(), + logger, + null, + null + ); + + // asyncBlobUpload calls onFailure with a wrapped CorruptFileException + Mockito.doAnswer(invocation -> { + ActionListener completionListener = invocation.getArgument(1); + completionListener.onFailure(new RuntimeException(new CorruptFileException("corrupted", "test_file"))); + return null; + }).when(asyncContainer).asyncBlobUpload(any(WriteContext.class), any()); + + Directory storeDirectory = newDirectory(); + String filename = "_100.si"; + IndexOutput indexOutput = storeDirectory.createOutput(filename, IOContext.DEFAULT); + indexOutput.writeString("data"); + CodecUtil.writeFooter(indexOutput); + indexOutput.close(); + storeDirectory.sync(List.of(filename)); + + CountDownLatch latch = new CountDownLatch(1); + AtomicReference failureRef = new AtomicReference<>(); + + boolean result = asyncDir.copyFrom( + storeDirectory, + filename, + filename + "__UUID", + IOContext.DEFAULT, + () -> {}, + new ActionListener<>() { + @Override + public void onResponse(Void unused) { + fail("Should not succeed"); + } + + @Override + public void onFailure(Exception e) { + failureRef.set(e); + latch.countDown(); + } + }, + false, + null + ); + + assertTrue(result); + assertTrue(latch.await(10, TimeUnit.SECONDS)); + assertTrue("Should be CorruptIndexException", failureRef.get() instanceof CorruptIndexException); + storeDirectory.close(); + } + + public void testCompletionListener_GenericException() throws Exception { + AsyncMultiStreamBlobContainer asyncContainer = mock(AsyncMultiStreamBlobContainer.class); + when(asyncContainer.remoteIntegrityCheckSupported()).thenReturn(false); + when(asyncContainer.path()).thenReturn(baseBlobPath); + + BlobStore asyncBlobStore = mock(BlobStore.class); + when(asyncBlobStore.blobContainer(baseBlobPath)).thenReturn(asyncContainer); + + DataFormatAwareRemoteDirectory asyncDir = new DataFormatAwareRemoteDirectory( + asyncBlobStore, + baseBlobPath, + UnaryOperator.identity(), + UnaryOperator.identity(), + UnaryOperator.identity(), + UnaryOperator.identity(), + new HashMap<>(), + logger, + null, + null + ); + + // asyncBlobUpload calls onFailure with a generic exception (not corrupt) + Mockito.doAnswer(invocation -> { + ActionListener completionListener = invocation.getArgument(1); + completionListener.onFailure(new IOException("network error")); + return null; + }).when(asyncContainer).asyncBlobUpload(any(WriteContext.class), any()); + + Directory storeDirectory = newDirectory(); + String filename = "_100.si"; + IndexOutput indexOutput = storeDirectory.createOutput(filename, IOContext.DEFAULT); + indexOutput.writeString("data"); + CodecUtil.writeFooter(indexOutput); + indexOutput.close(); + storeDirectory.sync(List.of(filename)); + + CountDownLatch latch = new CountDownLatch(1); + AtomicReference failureRef = new AtomicReference<>(); + + boolean result = asyncDir.copyFrom( + storeDirectory, + filename, + filename + "__UUID", + IOContext.DEFAULT, + () -> {}, + new ActionListener<>() { + @Override + public void onResponse(Void unused) { + fail("Should not succeed"); + } + + @Override + public void onFailure(Exception e) { + failureRef.set(e); + latch.countDown(); + } + }, + false, + null + ); + + assertTrue(result); + assertTrue(latch.await(10, TimeUnit.SECONDS)); + assertTrue("Should be IOException", failureRef.get() instanceof IOException); + storeDirectory.close(); + } + + // ═══════════════════════════════════════════════════════════════ + // Sync copyFrom with FileMetadata + // ═══════════════════════════════════════════════════════════════ + + public void testCopyFrom_FileMetadata_Sync() throws IOException { + DataFormatAwareStoreDirectory mockComposite = mock(DataFormatAwareStoreDirectory.class); + FileMetadata fm = new FileMetadata("lucene", "_0.cfs"); + + IndexInput mockInput = mock(IndexInput.class); + when(mockInput.length()).thenReturn(50L); + when(mockComposite.openInput(eq(fm.serialize()), eq(IOContext.DEFAULT))).thenReturn(mockInput); + + // Should write to base container since format is "lucene" + directory.copyFrom(mockComposite, fm.serialize(), "_0.cfs__UUID", IOContext.DEFAULT); + + verify(mockComposite).openInput(eq(fm.serialize()), eq(IOContext.DEFAULT)); + } + + // ═══════════════════════════════════════════════════════════════ + // Low priority upload path (content > 15GB triggers low priority) + // ═══════════════════════════════════════════════════════════════ + + public void testAsyncCopyFrom_LowPriorityUpload() throws Exception { + AsyncMultiStreamBlobContainer asyncContainer = mock(AsyncMultiStreamBlobContainer.class); + when(asyncContainer.remoteIntegrityCheckSupported()).thenReturn(false); + when(asyncContainer.path()).thenReturn(baseBlobPath); + + BlobStore asyncBlobStore = mock(BlobStore.class); + when(asyncBlobStore.blobContainer(baseBlobPath)).thenReturn(asyncContainer); + + DataFormatAwareRemoteDirectory asyncDir = new DataFormatAwareRemoteDirectory( + asyncBlobStore, + baseBlobPath, + UnaryOperator.identity(), + UnaryOperator.identity(), + UnaryOperator.identity(), + UnaryOperator.identity(), + new HashMap<>(), + logger, + null, + null + ); + + Mockito.doAnswer(invocation -> { + ActionListener completionListener = invocation.getArgument(1); + completionListener.onResponse(null); + return null; + }).when(asyncContainer).asyncBlobUpload(any(WriteContext.class), any()); + + Directory storeDirectory = newDirectory(); + String filename = "_100.si"; + IndexOutput indexOutput = storeDirectory.createOutput(filename, IOContext.DEFAULT); + indexOutput.writeString("data"); + CodecUtil.writeFooter(indexOutput); + indexOutput.close(); + storeDirectory.sync(List.of(filename)); + + CountDownLatch latch = new CountDownLatch(1); + + // Pass lowPriorityUpload=true + boolean result = asyncDir.copyFrom( + storeDirectory, + filename, + filename + "__UUID", + IOContext.DEFAULT, + () -> {}, + new ActionListener<>() { + @Override + public void onResponse(Void unused) { + latch.countDown(); + } + + @Override + public void onFailure(Exception e) { + fail("Should not fail: " + e.getMessage()); + } + }, + true, // lowPriorityUpload + null + ); + + assertTrue(result); + assertTrue(latch.await(10, TimeUnit.SECONDS)); + storeDirectory.close(); + } + + // ═══════════════════════════════════════════════════════════════ + // Blob Format Cache Tests + // ═══════════════════════════════════════════════════════════════ + + public void testReplaceBlobFormatCache() throws IOException { + // Initially no cache entries — deleteFile for a UUID key defaults to lucene (base container) + directory.deleteFile("_0.pqt__UUID1"); + verify(baseBlobContainer).deleteBlobsIgnoringIfNotExists(eq(Collections.singletonList("_0.pqt__UUID1"))); + + // Replace cache with parquet mapping + directory.getFormatBlobRouter().orElseThrow().replaceBlobFormatCache(Map.of("_0.pqt__UUID1", "parquet")); + directory.deleteFile("_0.pqt__UUID1"); + verify(parquetBlobContainer).deleteBlobsIgnoringIfNotExists(eq(Collections.singletonList("_0.pqt__UUID1"))); + } + + public void testUnregisterBlobFormat() throws IOException { + directory.getFormatBlobRouter().orElseThrow().registerBlobFormat("_0.pqt__UUID1", "parquet"); + directory.deleteFile("_0.pqt__UUID1"); + verify(parquetBlobContainer).deleteBlobsIgnoringIfNotExists(eq(Collections.singletonList("_0.pqt__UUID1"))); + + // Unregister — should fall back to lucene + directory.getFormatBlobRouter().orElseThrow().unregisterBlobFormat("_0.pqt__UUID1"); + directory.deleteFile("_0.pqt__UUID1"); + verify(baseBlobContainer).deleteBlobsIgnoringIfNotExists(eq(Collections.singletonList("_0.pqt__UUID1"))); + } + + public void testResolveFormat_CacheMiss_DefaultsToLucene() throws IOException { + // UUID-suffixed key not in cache — should warn and default to lucene + directory.deleteFile("_0.pqt__UUID_MISSING"); + verify(baseBlobContainer).deleteBlobsIgnoringIfNotExists(eq(Collections.singletonList("_0.pqt__UUID_MISSING"))); + } + + // ═══════════════════════════════════════════════════════════════ + // openBlockInput Tests + // ═══════════════════════════════════════════════════════════════ + + public void testOpenBlockInput_LuceneFile() throws IOException { + byte[] data = new byte[] { 1, 2, 3, 4, 5 }; + InputStream stream = new ByteArrayInputStream(data); + when(baseBlobContainer.readBlob("_0.cfe__UUID1", 0, 5)).thenReturn(stream); + + IndexInput input = directory.openBlockInput("_0.cfe__UUID1", 0, 5, 5, IOContext.DEFAULT); + assertNotNull(input); + input.close(); + } + + public void testOpenBlockInput_ParquetFile_WithCache() throws IOException { + directory.getFormatBlobRouter().orElseThrow().registerBlobFormat("_0.pqt__UUID1", "parquet"); + byte[] data = new byte[] { 10, 20, 30 }; + InputStream stream = new ByteArrayInputStream(data); + when(parquetBlobContainer.readBlob("_0.pqt__UUID1", 2, 3)).thenReturn(stream); + + IndexInput input = directory.openBlockInput("_0.pqt__UUID1", 2, 3, 10, IOContext.DEFAULT); + assertNotNull(input); + input.close(); + } + + public void testOpenBlockInput_InvalidPosition_Throws() { + expectThrows(IllegalArgumentException.class, () -> directory.openBlockInput("_0.cfe__UUID1", -1, 5, 10, IOContext.DEFAULT)); + } + + public void testOpenBlockInput_LengthExceedsFileLength_Throws() { + expectThrows(IllegalArgumentException.class, () -> directory.openBlockInput("_0.cfe__UUID1", 5, 10, 10, IOContext.DEFAULT)); + } +} diff --git a/server/src/test/java/org/opensearch/index/store/remote/FormatBlobRouterTests.java b/server/src/test/java/org/opensearch/index/store/remote/FormatBlobRouterTests.java new file mode 100644 index 0000000000000..04e5942b66dd4 --- /dev/null +++ b/server/src/test/java/org/opensearch/index/store/remote/FormatBlobRouterTests.java @@ -0,0 +1,210 @@ +/* + * 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.store.remote; + +import org.opensearch.common.blobstore.BlobContainer; +import org.opensearch.common.blobstore.BlobPath; +import org.opensearch.common.blobstore.BlobStore; +import org.opensearch.test.OpenSearchTestCase; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Tests for {@link FormatBlobRouter}. + */ +public class FormatBlobRouterTests extends OpenSearchTestCase { + + private BlobStore blobStore; + private BlobPath basePath; + private BlobContainer baseContainer; + + @Override + public void setUp() throws Exception { + super.setUp(); + blobStore = mock(BlobStore.class); + basePath = new BlobPath().add("indices").add("shard0").add("segments"); + baseContainer = mock(BlobContainer.class); + when(blobStore.blobContainer(basePath)).thenReturn(baseContainer); + } + + public void testLuceneFormatReturnsBaseContainer() { + FormatBlobRouter router = new FormatBlobRouter(blobStore, basePath); + assertSame(baseContainer, router.containerFor("lucene")); + } + + public void testNullFormatReturnsBaseContainer() { + FormatBlobRouter router = new FormatBlobRouter(blobStore, basePath); + assertSame(baseContainer, router.containerFor(null)); + } + + public void testMetadataFormatReturnsBaseContainer() { + FormatBlobRouter router = new FormatBlobRouter(blobStore, basePath); + assertSame(baseContainer, router.containerFor("metadata")); + } + + public void testNonLuceneFormatCreatesSubPathContainer() { + BlobContainer parquetContainer = mock(BlobContainer.class); + BlobPath parquetPath = basePath.add("parquet"); + when(blobStore.blobContainer(parquetPath)).thenReturn(parquetContainer); + + FormatBlobRouter router = new FormatBlobRouter(blobStore, basePath); + BlobContainer result = router.containerFor("parquet"); + + assertSame(parquetContainer, result); + verify(blobStore).blobContainer(parquetPath); + } + + public void testSameFormatReturnsSameContainer() { + BlobContainer parquetContainer = mock(BlobContainer.class); + when(blobStore.blobContainer(basePath.add("parquet"))).thenReturn(parquetContainer); + + FormatBlobRouter router = new FormatBlobRouter(blobStore, basePath); + BlobContainer first = router.containerFor("parquet"); + BlobContainer second = router.containerFor("parquet"); + + assertSame(first, second); + // blobContainer called once for basePath (constructor) + once for parquet (first access) + verify(blobStore, times(1)).blobContainer(basePath.add("parquet")); + } + + public void testBaseContainerAccessor() { + FormatBlobRouter router = new FormatBlobRouter(blobStore, basePath); + assertSame(baseContainer, router.baseContainer()); + } + + public void testRegisteredFormatsIncludesLuceneByDefault() { + FormatBlobRouter router = new FormatBlobRouter(blobStore, basePath); + assertTrue(router.registeredFormats().contains("lucene")); + assertEquals(1, router.registeredFormats().size()); + } + + public void testRegisteredFormatsGrowsWithAccess() { + when(blobStore.blobContainer(any(BlobPath.class))).thenReturn(mock(BlobContainer.class)); + + FormatBlobRouter router = new FormatBlobRouter(blobStore, basePath); + router.containerFor("parquet"); + router.containerFor("arrow"); + + assertTrue(router.registeredFormats().contains("lucene")); + assertTrue(router.registeredFormats().contains("parquet")); + assertTrue(router.registeredFormats().contains("arrow")); + assertEquals(3, router.registeredFormats().size()); + } + + public void testRegisterFormatPreCreatesContainer() { + BlobContainer parquetContainer = mock(BlobContainer.class); + when(blobStore.blobContainer(basePath.add("parquet"))).thenReturn(parquetContainer); + + FormatBlobRouter router = new FormatBlobRouter(blobStore, basePath); + router.registerFormat("parquet"); + + assertTrue(router.registeredFormats().contains("parquet")); + verify(blobStore).blobContainer(basePath.add("parquet")); + } + + public void testRegisterFormatIgnoresBasePathFormats() { + FormatBlobRouter router = new FormatBlobRouter(blobStore, basePath); + router.registerFormat("lucene"); + router.registerFormat("metadata"); + router.registerFormat(null); + + // Only lucene should be registered (default), no extra containers created + assertEquals(1, router.registeredFormats().size()); + } + + public void testFormatNameIsLowercasedForPath() { + BlobContainer container = mock(BlobContainer.class); + when(blobStore.blobContainer(basePath.add("parquet"))).thenReturn(container); + + FormatBlobRouter router = new FormatBlobRouter(blobStore, basePath); + router.containerFor("Parquet"); + + // Should create path with lowercase "parquet" + verify(blobStore).blobContainer(basePath.add("parquet")); + } + + // ═══════════════════════════════════════════════════════════════ + // Blob Format Cache Tests + // ═══════════════════════════════════════════════════════════════ + + public void testResolveFormat_DefaultsToLucene() { + FormatBlobRouter router = new FormatBlobRouter(blobStore, basePath); + assertEquals("lucene", router.resolveFormat("_0.cfs__UUID")); + } + + public void testRegisterBlobFormat_ThenResolve() { + FormatBlobRouter router = new FormatBlobRouter(blobStore, basePath); + router.registerBlobFormat("_0.pqt__UUID", "parquet"); + assertEquals("parquet", router.resolveFormat("_0.pqt__UUID")); + } + + public void testRegisterBlobFormat_NullKeyIgnored() { + FormatBlobRouter router = new FormatBlobRouter(blobStore, basePath); + router.registerBlobFormat(null, "parquet"); + // Should not throw, cache unchanged + assertEquals("lucene", router.resolveFormat("anything")); + } + + public void testRegisterBlobFormat_NullFormatIgnored() { + FormatBlobRouter router = new FormatBlobRouter(blobStore, basePath); + router.registerBlobFormat("_0.pqt__UUID", null); + assertEquals("lucene", router.resolveFormat("_0.pqt__UUID")); + } + + public void testUnregisterBlobFormat() { + FormatBlobRouter router = new FormatBlobRouter(blobStore, basePath); + router.registerBlobFormat("_0.pqt__UUID", "parquet"); + assertEquals("parquet", router.resolveFormat("_0.pqt__UUID")); + + router.unregisterBlobFormat("_0.pqt__UUID"); + assertEquals("lucene", router.resolveFormat("_0.pqt__UUID")); + } + + public void testUnregisterBlobFormat_NullKeyIgnored() { + FormatBlobRouter router = new FormatBlobRouter(blobStore, basePath); + router.registerBlobFormat("_0.pqt__UUID", "parquet"); + router.unregisterBlobFormat(null); + // Cache unchanged + assertEquals("parquet", router.resolveFormat("_0.pqt__UUID")); + } + + public void testReplaceBlobFormatCache() { + FormatBlobRouter router = new FormatBlobRouter(blobStore, basePath); + router.registerBlobFormat("old_key", "parquet"); + + router.replaceBlobFormatCache(java.util.Map.of("new_key", "arrow")); + + assertEquals("lucene", router.resolveFormat("old_key")); + assertEquals("arrow", router.resolveFormat("new_key")); + } + + public void testClearBlobFormatCache() { + FormatBlobRouter router = new FormatBlobRouter(blobStore, basePath); + router.registerBlobFormat("_0.pqt__UUID", "parquet"); + assertEquals("parquet", router.resolveFormat("_0.pqt__UUID")); + + router.clearBlobFormatCache(); + assertEquals("lucene", router.resolveFormat("_0.pqt__UUID")); + } + + public void testReplaceBlobFormatCache_ImmutableSnapshot() { + FormatBlobRouter router = new FormatBlobRouter(blobStore, basePath); + java.util.Map mutable = new java.util.HashMap<>(); + mutable.put("key1", "parquet"); + router.replaceBlobFormatCache(mutable); + + // Mutating the original map should not affect the cache + mutable.put("key2", "arrow"); + assertEquals("lucene", router.resolveFormat("key2")); + } +}