diff --git a/packaging/pom.xml b/packaging/pom.xml index df4d33309e31..54566635f81c 100644 --- a/packaging/pom.xml +++ b/packaging/pom.xml @@ -453,6 +453,11 @@ hive-standalone-metastore-rest-catalog ${project.version} + + org.apache.hive + hive-standalone-metastore-search + ${project.version} + org.apache.hadoop hadoop-hdfs-client diff --git a/packaging/src/main/assembly/src.xml b/packaging/src/main/assembly/src.xml index 4dc658e1871e..5f1fc9d149a1 100644 --- a/packaging/src/main/assembly/src.xml +++ b/packaging/src/main/assembly/src.xml @@ -106,6 +106,7 @@ storage-api/**/* standalone-metastore/metastore-common/**/* standalone-metastore/metastore-client/**/* + standalone-metastore/metastore-search/**/* standalone-metastore/metastore-server/**/* standalone-metastore/metastore-tools/**/* standalone-metastore/metastore-rest-catalog/**/* diff --git a/standalone-metastore/metastore-search/pom.xml b/standalone-metastore/metastore-search/pom.xml new file mode 100644 index 000000000000..19463414d174 --- /dev/null +++ b/standalone-metastore/metastore-search/pom.xml @@ -0,0 +1,219 @@ + + + + + hive-standalone-metastore + org.apache.hive + 4.3.0-SNAPSHOT + + 4.0.0 + hive-standalone-metastore-search + Hive Metastore Search + + .. + 1.25.1 + 0.36.0 + 10.5.0 + + + + org.apache.hive + hive-standalone-metastore-server + ${hive.version} + + + org.apache.hadoop + hadoop-common + + + org.slf4j + slf4j-log4j12 + + + org.slf4j + slf4j-reload4j + + + ch.qos.reload4j + reload4j + + + commons-logging + commons-logging + + + + + org.apache.hadoop + hadoop-hdfs-client + + + org.slf4j + slf4j-log4j12 + + + org.slf4j + slf4j-reload4j + + + ch.qos.reload4j + reload4j + + + commons-logging + commons-logging + + + + + org.apache.commons + commons-lang3 + + + com.google.guava + guava + + + com.fasterxml.jackson.core + jackson-databind + + + com.microsoft.onnxruntime + onnxruntime + ${onnxruntime.version} + + + ai.djl.huggingface + tokenizers + ${djl.tokenizers.version} + + + org.apache.lucene + lucene-core + ${lucene.version} + + + org.apache.lucene + lucene-analysis-common + ${lucene.version} + + + org.apache.lucene + lucene-queryparser + ${lucene.version} + + + org.slf4j + slf4j-api + + + junit + junit + test + + + org.apache.hive + hive-standalone-metastore-server + ${hive.version} + tests + test + + + org.mockito + mockito-core + test + + + org.apache.derby + derby + test + + + org.apache.hive.hcatalog + hive-hcatalog-server-extensions + ${hive.version} + test + + + org.apache.logging.log4j + log4j-slf4j-impl + test + + + + + + org.apache.maven.plugins + maven-antrun-plugin + + + setup-test-dirs + process-test-resources + + run + + + + + + + + + + + + + + setup-metastore-scripts + process-test-resources + + run + + + + + + + + + + + + + + org.apache.maven.plugins + maven-surefire-plugin + + true + false + -Xmx2048m --add-opens java.base/java.lang=ALL-UNNAMED --add-opens java.base/java.lang.reflect=ALL-UNNAMED --add-opens java.base/java.util.concurrent.atomic=ALL-UNNAMED --add-opens java.base/java.net=ALL-UNNAMED --add-opens java.base/java.util=ALL-UNNAMED --add-opens java.base/java.nio=ALL-UNNAMED --add-opens java.base/java.util.concurrent=ALL-UNNAMED --add-opens java.base/java.util.regex=ALL-UNNAMED --add-opens java.sql/java.sql=ALL-UNNAMED --add-opens java.base/java.io=ALL-UNNAMED --add-opens java.base/java.text=ALL-UNNAMED + + ${project.build.directory} + true + ${derby.version} + ${test.tmp.dir}/derby.log + ${test.tmp.dir} + jdbc:derby:memory:${test.tmp.dir}/junit_metastore_db;create=true + false + ${test.tmp.dir} + ${test.warehouse.scheme}${test.warehouse.dir} + ${test.warehouse.scheme}${test.warehouse.external.dir} + + + ${project.build.testOutputDirectory} + + + + + + diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hadoop/hive/metastore/handler/SearchProvider.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hadoop/hive/metastore/handler/SearchProvider.java new file mode 100644 index 000000000000..18a90359d43b --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hadoop/hive/metastore/handler/SearchProvider.java @@ -0,0 +1,117 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hadoop.hive.metastore.handler; + +import java.io.IOException; +import java.util.Objects; +import java.util.concurrent.atomic.AtomicReference; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hive.search.exception.IndexNotReadyException; +import org.apache.hive.search.exception.InitializeException; +import org.apache.hive.search.exception.SearchException; +import org.apache.hive.search.search.LuceneSearchBackend; +import org.apache.hive.search.search.SearchBackend; +import org.apache.hive.search.search.SearchQuery; +import org.apache.hive.search.search.TableSearchResult; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Single server-side provider that owns the lifecycle of the {@link SearchBackend}. */ +public final class SearchProvider implements AutoCloseable { + private static final Logger LOG = LoggerFactory.getLogger(SearchProvider.class); + private static final AtomicReference INSTANCE = new AtomicReference<>(); + + private final Configuration configuration; + private final SearchBackend backend; + + private SearchProvider(Configuration configuration, SearchBackend backend) { + this.configuration = configuration; + this.backend = backend; + } + + /** Creates and initializes a provider with the supplied backend implementation. */ + private static SearchProvider create(Configuration configuration, SearchBackend backend) + throws InitializeException, IOException { + Objects.requireNonNull(configuration, "configuration"); + Objects.requireNonNull(backend, "backend"); + backend.initialize(configuration); + LOG.info("Installed search backend {}", backend.getClass().getSimpleName()); + return new SearchProvider(configuration, backend); + } + + public static SearchProvider install(Configuration configuration) + throws InitializeException, IOException { + return install(configuration, new LuceneSearchBackend()); + } + + /** Installs the process-wide provider using the supplied backend implementation. */ + public static SearchProvider install(Configuration configuration, SearchBackend backend) + throws InitializeException, IOException { + SearchProvider provider; + if ((provider = INSTANCE.get()) != null) { + return provider; + } + synchronized (SearchProvider.class) { + if ((provider = INSTANCE.get()) == null) { + provider = create(configuration, backend); + INSTANCE.getAndSet(provider); + } + } + return provider; + } + + /** Returns the installed provider. */ + public static SearchProvider get() { + SearchProvider provider = INSTANCE.get(); + if (provider == null) { + throw new IllegalStateException("SearchProvider is not installed"); + } + return provider; + } + + /** Clears the installed provider; intended for tests. */ + public static void reset() throws Exception { + SearchProvider previous = INSTANCE.getAndSet(null); + if (previous != null) { + previous.close(); + } + } + + public Configuration configuration() { + return configuration; + } + + public SearchBackend backend() { + return backend; + } + + public boolean isReady() throws IndexNotReadyException { + return backend.isReady(); + } + + public TableSearchResult search(SearchQuery query) + throws SearchException, InitializeException, IOException { + return backend.search(query); + } + + @Override + public void close() throws Exception { + backend.close(); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/config/IndexConfig.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/config/IndexConfig.java new file mode 100644 index 000000000000..996683b838c2 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/config/IndexConfig.java @@ -0,0 +1,156 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.config; + +import java.time.Duration; + +import org.apache.commons.lang3.StringUtils; +import org.apache.hadoop.conf.Configuration; + +public record IndexConfig(Configuration configuration) { + public static final String INDEX_RAM_SIZE = "metastore.index.write.ram.size"; + public static final int INDEX_RAM_SIZE_DEFAULT = 128 * 1024 * 1024; + + public static final String FLUSH_INTERVAL_MS = "metastore.index.flush.interval.ms"; + public static final long FLUSH_INTERVAL_MS_DEFAULT = 60 * 1000L; + + public static final String INDEX_NAME = "metastore.index.name"; + public static final String INDEX_NAME_DEFAULT = "hive_tables"; + + public static final String POLL_NOTIFICATION_INTERVAL = "metastore.index.poll.notification.ms"; + public static final long POLL_NOTIFICATION_INTERVAL_DEFAULT = 1000; + + public static final String INDEX_SYNC_INTERVAL = "metastore.index.sync.interval.minutes"; + public static final long INDEX_SYNC_INTERVAL_DEFAULT = 360; + + /** Force a metadata-only commit when this many events are ahead of the last committed checkpoint. */ + public static final String FORCE_FLUSH_EVENT_GAP = "metastore.index.force.flush.event.gap"; + public static final long FORCE_FLUSH_EVENT_GAP_DEFAULT = 3000L; + + /** Tables per metastore fetch batch during leader bootstrap. */ + public static final String BOOTSTRAP_BATCH_SIZE = "metastore.index.bootstrap.batch.size"; + public static final int BOOTSTRAP_BATCH_SIZE_DEFAULT = 2000; + + /** In-flight document batches between fetch workers and the index writer. */ + public static final String BOOTSTRAP_QUEUE_DEPTH = "metastore.index.bootstrap.queue.depth"; + public static final int BOOTSTRAP_QUEUE_DEPTH_DEFAULT = 16; + + /** Parallel metastore fetch workers during bootstrap. */ + public static final String BOOTSTRAP_FETCH_THREADS = "metastore.index.bootstrap.fetch.threads"; + public static final int BOOTSTRAP_FETCH_THREADS_DEFAULT = 4; + + /** Lucene commits after this many Lucene auto-flushes during incremental indexing. 0 commits immediately. */ + public static final String COMMIT_FLUSHES = "metastore.index.commit.flushes"; + public static final int COMMIT_FLUSHES_DEFAULT = 3; + + /** Progress log interval during bootstrap. */ + public static final String BOOTSTRAP_PROGRESS_INTERVAL_MS = + "metastore.index.bootstrap.progress.interval.ms"; + public static final long BOOTSTRAP_PROGRESS_INTERVAL_MS_DEFAULT = 30_000L; + + /** Max HMS events fetched per poll. */ + public static final String EVENT_POLL_MAX = "metastore.index.event.poll.max"; + public static final int EVENT_POLL_MAX_DEFAULT = 3000; + + /** Consecutive batch failures before falling back to single-event apply. */ + public static final String EVENT_BATCH_MAX_FAILURES = + "metastore.index.event.batch.max.failures"; + public static final int EVENT_BATCH_MAX_FAILURES_DEFAULT = 3; + + /** Extra sleep after a failed batch before the poller retries (ms). */ + public static final String EVENT_FAILURE_BACKOFF_MS = "metastore.index.event.failure.backoff.ms"; + public static final long EVENT_FAILURE_BACKOFF_MS_DEFAULT = 10_000L; + + /** Skip a single poison event during individual fallback (index may diverge for that table). */ + public static final String EVENT_SKIP_POISON = "metastore.index.event.skip.poison"; + public static final boolean EVENT_SKIP_POISON_DEFAULT = false; + + /** Poller sleep after index is marked unhealthy (ms). */ + public static final String EVENT_UNHEALTHY_BACKOFF_MS = "metastore.index.event.unhealthy.backoff.ms"; + public static final long EVENT_UNHEALTHY_BACKOFF_MS_DEFAULT = 10 * 60 * 1000L; + + public int getWriteBufferSize() { + return configuration.getInt(INDEX_RAM_SIZE, INDEX_RAM_SIZE_DEFAULT); + } + + public Duration getFlushInterval() { + return Duration.ofMillis(configuration.getLong(FLUSH_INTERVAL_MS, FLUSH_INTERVAL_MS_DEFAULT)); + } + + public double getWriteBufferSizeMb() { + return getWriteBufferSize() / (1024.0 * 1024.0); + } + + public long getPollNotificationInterval() { + return configuration.getLong(POLL_NOTIFICATION_INTERVAL, POLL_NOTIFICATION_INTERVAL_DEFAULT); + } + + public long getSyncInterval() { + return configuration.getLong(INDEX_SYNC_INTERVAL, INDEX_SYNC_INTERVAL_DEFAULT) * 60 * 1000; + } + + public long getForceFlushEventGap() { + return configuration.getLong(FORCE_FLUSH_EVENT_GAP, FORCE_FLUSH_EVENT_GAP_DEFAULT); + } + + public int getBootstrapBatchSize() { + return configuration.getInt(BOOTSTRAP_BATCH_SIZE, BOOTSTRAP_BATCH_SIZE_DEFAULT); + } + + public int getBootstrapQueueDepth() { + return configuration.getInt(BOOTSTRAP_QUEUE_DEPTH, BOOTSTRAP_QUEUE_DEPTH_DEFAULT); + } + + public int getBootstrapFetchThreads() { + return configuration.getInt(BOOTSTRAP_FETCH_THREADS, BOOTSTRAP_FETCH_THREADS_DEFAULT); + } + + public int getCommitFlushes() { + return configuration.getInt(COMMIT_FLUSHES, COMMIT_FLUSHES_DEFAULT); + } + + public long getBootstrapProgressIntervalMs() { + return configuration.getLong(BOOTSTRAP_PROGRESS_INTERVAL_MS, + BOOTSTRAP_PROGRESS_INTERVAL_MS_DEFAULT); + } + + public int getEventPollMax() { + return configuration.getInt(EVENT_POLL_MAX, EVENT_POLL_MAX_DEFAULT); + } + + public int getEventBatchMaxFailures() { + return configuration.getInt(EVENT_BATCH_MAX_FAILURES, EVENT_BATCH_MAX_FAILURES_DEFAULT); + } + + public long getEventFailureBackoffMs() { + return configuration.getLong(EVENT_FAILURE_BACKOFF_MS, EVENT_FAILURE_BACKOFF_MS_DEFAULT); + } + + public boolean isEventSkipPoison() { + return configuration.getBoolean(EVENT_SKIP_POISON, EVENT_SKIP_POISON_DEFAULT); + } + + public long getEventUnhealthyBackoffMs() { + return configuration.getLong(EVENT_UNHEALTHY_BACKOFF_MS, EVENT_UNHEALTHY_BACKOFF_MS_DEFAULT); + } + + public String indexName() { + String name = configuration.get(INDEX_NAME); + return StringUtils.isEmpty(name) ? INDEX_NAME_DEFAULT : name; + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/config/IndexStoreConfig.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/config/IndexStoreConfig.java new file mode 100644 index 000000000000..9bb448aecd48 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/config/IndexStoreConfig.java @@ -0,0 +1,64 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.config; + +import java.net.URI; +import java.nio.file.Path; +import java.nio.file.Paths; + +import org.apache.commons.lang3.StringUtils; +import org.apache.hadoop.conf.Configuration; +import org.apache.hive.search.exception.IndexIOException; + +public record IndexStoreConfig(Configuration configuration, String indexName) { + public static final String LOCAL_PATH = "metastore.index.local.path"; + public static final String REMOTE_URI = "metastore.index.backup.remote.uri"; + public static final String MEMORY = "metastore.index.use.memory"; + + public static final Path DEFAULT_WORKDIR = Paths.get(System.getProperty("user.dir"), "indexes"); + + public Path getLocalPath() { + String path = configuration.get(LOCAL_PATH, DEFAULT_WORKDIR.toString()); + return Path.of(path); + } + + public String getRemoteUri() { + return configuration.get(REMOTE_URI, ""); + } + + public boolean hasRemote() { + String uri = getRemoteUri(); + return StringUtils.isNotEmpty(uri); + } + + public boolean useMemory() { + return configuration.getBoolean(MEMORY, false); + } + + public boolean isDistributed() { + return hasRemote(); + } + + public static void validateRemoteUri(String uriText) throws IndexIOException { + try { + URI.create(uriText); + } catch (IllegalArgumentException e) { + throw new IndexIOException("invalid store.remote.uri: " + uriText, e); + } + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/config/InferenceConfig.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/config/InferenceConfig.java new file mode 100644 index 000000000000..02d5430e1bb9 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/config/InferenceConfig.java @@ -0,0 +1,147 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.config; + +import java.io.IOException; +import java.net.URI; + +import org.apache.commons.lang3.StringUtils; +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.FileStatus; +import org.apache.hadoop.fs.FileSystem; +import org.apache.hadoop.fs.Path; +import org.apache.hive.search.exception.InitializeException; +import org.apache.hive.search.inference.EmbeddingPrompt; + +public record InferenceConfig(Configuration configuration) { + public static final String MODEL_ONNX_FILE = "model.onnx"; + public static final String TOKENIZER = "tokenizer.json"; + + public static final String MODEL_LOCAL_DIR = "metastore.inference.local.dir"; + public static final String MODEL_LOCAL_DIR_DEFAULT = System.getProperty("java.io.tmpdir") + "/hivesearch-cache"; + public static final String MODEL_REMOTE_DIR = "metastore.inference.remote.dir"; + + public static final String MODEL_NAME = "metastore.inference.embedding.name"; + public static final String EMBEDDING_PROMPT_DOC = "metastore.inference.embedding.prompt.doc"; + public static final String EMBEDDING_PROMPT_QUERY = "metastore.inference.embedding.prompt.query"; + + public static final String EMBEDDING_CACHE_ENABLED = "metastore.inference.embedding.cache.enabled"; + public static final boolean EMBEDDING_CACHE_ENABLED_DEFAULT = true; + + public static final String EMBEDDING_CACHE_MAX_ENTRIES = "metastore.inference.embedding.cache.max.entries"; + public static final int EMBEDDING_CACHE_MAX_ENTRIES_DEFAULT = 100_000; + + public static final String EMBEDDING_THREADS = "metastore.inference.embedding.threads"; + + public static final String EMBEDDING_MAX_SEQ_LENGTH = "metastore.inference.embedding.max_seq_length"; + public static final int EMBEDDING_MAX_SEQ_LENGTH_DEFAULT = 512; + + static int defaultEmbeddingThreads() { + return Math.min(4, Math.max(1, Runtime.getRuntime().availableProcessors())); + } + + public EmbeddingModelSpec embedding() throws InitializeException, IOException { + String modelName = modelName(); + if (StringUtils.isEmpty(modelName)) { + throw new InitializeException("No model configured for embedding the search"); + } + URI modelPath = URI.create(configuration.get(MODEL_LOCAL_DIR, MODEL_LOCAL_DIR_DEFAULT)); + Path lPath = new Path("file://" + modelPath.getPath()); + if (requireConfigured(lPath, modelName)) { + String remotePath = configuration.get(MODEL_REMOTE_DIR); + String message = "Can't find the model or tokenizer in %s for " + modelName; + if (StringUtils.isEmpty(remotePath)) { + throw new InitializeException(String.format(message, lPath)); + } + Path rPath = new Path(remotePath); + if (requireConfigured(rPath, modelName)) { + throw new InitializeException(String.format(message, rPath)); + } + Path dest = new Path(lPath, modelName); + Path src = new Path(rPath, modelName); + lPath.getFileSystem(configuration).mkdirs(dest); + FileSystem fileSystem = rPath.getFileSystem(configuration); + fileSystem.copyToLocalFile(new Path(src, MODEL_ONNX_FILE), new Path(dest, MODEL_ONNX_FILE)); + fileSystem.copyToLocalFile(new Path(src, TOKENIZER), new Path(dest, TOKENIZER)); + } + EmbeddingPrompt prompt = new EmbeddingPrompt(configuration.get(EMBEDDING_PROMPT_DOC, "passage: "), + configuration.get(EMBEDDING_PROMPT_QUERY, "query: ")); + return new EmbeddingModelSpec(modelName, new Path(lPath, modelName), prompt); + } + + private boolean requireConfigured(Path modelPath, String modelName) { + try { + FileSystem fileSystem = modelPath.getFileSystem(configuration); + FileStatus[] fileStatuses = fileSystem.listStatus(new Path(modelPath, modelName), path -> { + String fileName = path.getName(); + return fileName.equals(MODEL_ONNX_FILE) || fileName.equals(TOKENIZER); + }); + if (fileStatuses == null || fileStatuses.length != 2) { + return true; + } + return fileStatuses[0].isDirectory() || fileStatuses[1].isDirectory(); + } catch (IOException e) { + return true; + } + } + + public String modelName() { + return configuration.get(MODEL_NAME); + } + + public boolean isEmbeddingCacheEnabled() { + return configuration.getBoolean(EMBEDDING_CACHE_ENABLED, EMBEDDING_CACHE_ENABLED_DEFAULT); + } + + public int getEmbeddingCacheMaxEntries() { + return configuration.getInt(EMBEDDING_CACHE_MAX_ENTRIES, EMBEDDING_CACHE_MAX_ENTRIES_DEFAULT); + } + + public int getEmbeddingThreads() { + return Math.max(1, configuration.getInt(EMBEDDING_THREADS, defaultEmbeddingThreads())); + } + + public int getEmbeddingMaxSeqLength() { + return Math.max(1, configuration.getInt(EMBEDDING_MAX_SEQ_LENGTH, EMBEDDING_MAX_SEQ_LENGTH_DEFAULT)); + } + + public static class EmbeddingModelSpec { + private final String model; + private final java.nio.file.Path modelDir; + private final EmbeddingPrompt prompt; + + public EmbeddingModelSpec(String model, Path path, EmbeddingPrompt promp) { + this.model = model; + this.prompt = promp == null ? + EmbeddingPrompt.none() : promp; + this.modelDir = java.nio.file.Path.of(path.toUri()); + } + + public String getModel() { + return model; + } + + public java.nio.file.Path getModelDir() { + return modelDir; + } + + public EmbeddingPrompt getPrompt() { + return prompt; + } + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/config/SearchConfig.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/config/SearchConfig.java new file mode 100644 index 000000000000..1b05c69f2ea2 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/config/SearchConfig.java @@ -0,0 +1,91 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.config; + +import java.time.Duration; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hive.search.exception.SearchException; + +public record SearchConfig(Configuration configuration) { + public static final String REFRESH_INTERVAL_SECONDS = "metastore.search.refresh.interval.seconds"; + public static final long REFRESH_INTERVAL_SECONDS_DEFAULT = 1L; + + public static final String DEFAULT_LIMIT = "metastore.search.default.limit"; + public static final int DEFAULT_LIMIT_DEFAULT = 10; + + public static final String INIT_READY_TIMEOUT_MS = "metastore.search.init.ready.timeout.ms"; + public static final long INIT_READY_TIMEOUT_MS_DEFAULT = 300L; + + public static final String HYBRID_SEMANTIC_WEIGHT = "metastore.search.hybrid.semantic.weight"; + public static final float HYBRID_SEMANTIC_WEIGHT_DEFAULT = 0.4f; + + public static final String FUSION_PRIOR = "metastore.search.fusion.prior"; + public static final float FUSION_PRIOR_DEFAULT = 0.5f; + + public static final String BAYESIAN_SAMPLES = "metastore.search.bayesian.samples"; + public static final int BAYESIAN_SAMPLES_DEFAULT = 100; + + public static final String BAYESIAN_TOKENS_PER_QUERY = "metastore.search.bayesian.tokens.per.query"; + public static final int BAYESIAN_TOKENS_PER_QUERY_DEFAULT = 5; + + public static final String BAYESIAN_SEED = "metastore.search.bayesian.seed"; + public static final long BAYESIAN_SEED_DEFAULT = 42L; + + public Duration getRefreshInterval() { + return Duration.ofSeconds( + configuration.getLong(REFRESH_INTERVAL_SECONDS, REFRESH_INTERVAL_SECONDS_DEFAULT)); + } + + public int getDefaultLimit() { + return configuration.getInt(DEFAULT_LIMIT, DEFAULT_LIMIT_DEFAULT); + } + + public long getInitReadyTimeoutMs() { + return configuration.getLong(INIT_READY_TIMEOUT_MS, INIT_READY_TIMEOUT_MS_DEFAULT); + } + + public float getHybridSemanticWeight() throws SearchException { + float weight = configuration.getFloat(HYBRID_SEMANTIC_WEIGHT, HYBRID_SEMANTIC_WEIGHT_DEFAULT); + if (weight >= 1.0f || weight <= 0.0f) { + throw new SearchException("Invalid hybrid semantic weight, " + + "it must be in (0, 1), but got " + weight); + } + return weight; + } + + public float getHybridMatchWeight() throws SearchException { + return 1.0f - getHybridSemanticWeight(); + } + + public float getFusionPrior() { + return configuration.getFloat(FUSION_PRIOR, FUSION_PRIOR_DEFAULT); + } + + public int getBayesianSamples() { + return configuration.getInt(BAYESIAN_SAMPLES, BAYESIAN_SAMPLES_DEFAULT); + } + + public int getBayesianTokensPerQuery() { + return configuration.getInt(BAYESIAN_TOKENS_PER_QUERY, BAYESIAN_TOKENS_PER_QUERY_DEFAULT); + } + + public long getBayesianSeed() { + return configuration.getLong(BAYESIAN_SEED, BAYESIAN_SEED_DEFAULT); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/exception/IndexException.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/exception/IndexException.java new file mode 100644 index 000000000000..24a5dd13f801 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/exception/IndexException.java @@ -0,0 +1,36 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.exception; + +/** Failure while building, updating, or embedding content for the search index. */ +public class IndexException extends Exception { + public IndexException(String message) { + super(message); + } + + public IndexException(String message, Throwable cause) { + super(message, cause); + } + + public static IndexException wrap(String message, Throwable cause) { + if (cause instanceof IndexException indexException) { + return indexException; + } + return new IndexException(message, cause); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/exception/IndexIOException.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/exception/IndexIOException.java new file mode 100644 index 000000000000..dcb7685805b6 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/exception/IndexIOException.java @@ -0,0 +1,38 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.exception; + +import java.io.IOException; + +/** Failure while reading, writing, syncing, or restoring index files and manifests. */ +public class IndexIOException extends IOException { + public IndexIOException(String message) { + super(message); + } + + public IndexIOException(String message, Throwable cause) { + super(message, cause); + } + + public static IndexIOException wrap(Exception cause) { + if (cause instanceof IndexIOException indexIOException) { + return indexIOException; + } + return new IndexIOException(cause.getMessage(), cause); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/exception/IndexNotHealthyException.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/exception/IndexNotHealthyException.java new file mode 100644 index 000000000000..da1d5f2f5689 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/exception/IndexNotHealthyException.java @@ -0,0 +1,30 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.exception; + +import java.io.IOException; + +public class IndexNotHealthyException extends IOException { + public IndexNotHealthyException(String message) { + super(message); + } + + public IndexNotHealthyException(String message, Throwable cause) { + super(message, cause); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/exception/IndexNotReadyException.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/exception/IndexNotReadyException.java new file mode 100644 index 000000000000..518bb42d65df --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/exception/IndexNotReadyException.java @@ -0,0 +1,30 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.exception; + +import java.io.IOException; + +public class IndexNotReadyException extends IOException { + public IndexNotReadyException(String message) { + super(message); + } + + public IndexNotReadyException(String message, Throwable cause) { + super(message, cause); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/exception/InitializeException.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/exception/InitializeException.java new file mode 100644 index 000000000000..77e935b5e1c1 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/exception/InitializeException.java @@ -0,0 +1,36 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.exception; + +/** Failure while opening or bootstrapping the search service and its dependencies. */ +public class InitializeException extends Exception { + public InitializeException(String message) { + super(message); + } + + public InitializeException(String message, Throwable cause) { + super(message, cause); + } + + public static InitializeException wrap(String message, Throwable cause) { + if (cause instanceof InitializeException initializeException) { + return initializeException; + } + return new InitializeException(message, cause); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/exception/SearchException.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/exception/SearchException.java new file mode 100644 index 000000000000..e68e1a929a7b --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/exception/SearchException.java @@ -0,0 +1,29 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.exception; + +/** Failure while executing or preparing a table search request. */ +public class SearchException extends Exception { + public SearchException(String message) { + super(message); + } + + public SearchException(String message, Throwable cause) { + super(message, cause); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/FlushTrackingWriter.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/FlushTrackingWriter.java new file mode 100644 index 000000000000..ff3b13620f01 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/FlushTrackingWriter.java @@ -0,0 +1,48 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.index; + +import java.io.IOException; + +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.store.Directory; + +/** Tracks Lucene auto-flushes so callers can gate commits separately. */ +final class FlushTrackingWriter extends IndexWriter { + private int flushesSinceCommit; + + FlushTrackingWriter(Directory directory, IndexWriterConfig config) + throws IOException { + super(directory, config); + } + + int flushesSinceCommit() { + return flushesSinceCommit; + } + + void resetFlushTracking() { + flushesSinceCommit = 0; + } + + @Override + protected void doAfterFlush() throws IOException { + super.doAfterFlush(); + flushesSinceCommit++; + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/IndexManager.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/IndexManager.java new file mode 100644 index 000000000000..83bcd0bf031b --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/IndexManager.java @@ -0,0 +1,199 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.index; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Optional; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hive.search.exception.IndexException; +import org.apache.hive.search.mapping.IndexMapping; +import org.apache.hive.search.exception.IndexNotHealthyException; +import org.apache.hive.search.index.manifest.IndexManifest; +import org.apache.hive.search.exception.IndexIOException; +import org.apache.hive.search.index.store.LocalStateClient; +import org.apache.hive.search.index.store.IndexBackupUtils; +import org.apache.hive.search.index.store.IndexStateClient; +import org.apache.hive.search.config.IndexStoreConfig; +import org.apache.hive.search.metastore.MetastoreEventListener; +import org.apache.lucene.store.ByteBuffersDirectory; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.MMapDirectory; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class IndexManager implements AutoCloseable, MetastoreEventListener { + private static final Logger LOG = LoggerFactory.getLogger(IndexManager.class); + private final IndexMapping mapping; + private final Directory directory; + private final IndexStateClient localIndex; + private final IndexStateClient remoteIndex; + private volatile IndexNotHealthyException exception; + private volatile long indexedNid; + + public IndexManager(IndexMapping mapping, Directory directory, IndexStateClient localIndex, + IndexStateClient remoteIndex) { + this.mapping = mapping; + this.directory = directory; + this.localIndex = localIndex; + this.remoteIndex = remoteIndex; + } + + public static IndexManager open(IndexMapping mapping, Configuration conf) throws IOException { + IndexStoreConfig store = mapping.store(); + Directory directory; + if (store.useMemory()) { + directory = new ByteBuffersDirectory(); + } else { + directory = openDiskDirectory(mapping.indexName(), store.getLocalPath()); + } + + IndexStateClient local = + new LocalStateClient(directory, mapping.indexName()); + IndexStateClient remote = null; + if (store.hasRemote()) { + remote = IndexBackupUtils.openRemote(store.getRemoteUri(), mapping.indexName(), conf); + } + return new IndexManager(mapping, directory, local, remote); + } + + private static Directory openDiskDirectory(String indexName, Path basePath) throws IOException { + Path indexPath = basePath.resolve(indexName); + Files.createDirectories(indexPath); + return new MMapDirectory(indexPath); + } + + public boolean hasBackup() { + return remoteIndex != null; + } + + public boolean syncBackup() throws IOException { + if (remoteIndex == null) { + return false; + } + if (localIndex.readManifest().isEmpty()) { + return false; + } + return IndexBackupUtils.syncToBackup(localIndex, remoteIndex); + } + + public boolean restoreBackup() { + try { + if (remoteIndex == null) { + return false; + } + return IndexBackupUtils.restoreFromBackup(localIndex, remoteIndex); + } catch (IOException e) { + LOG.warn("Cannot restore the index from remote backup", e); + return false; + } + } + + public void resolveInterruptedRestore() { + try { + IndexBackupUtils.resolveInterruptedRestore(localIndex); + } catch (IOException e) { + LOG.warn("Cannot resolve interrupted local index restore", e); + } + } + + public Optional readLocalManifest() { + try { + return localIndex.readManifest(); + } catch (IOException e) { + LOG.warn("Cannot read the local index manifest", e); + return Optional.empty(); + } + } + + public Optional readRemoteManifest() { + try { + if (remoteIndex == null) { + return Optional.empty(); + } + return remoteIndex.readManifest(); + } catch (IOException e) { + LOG.warn("Cannot read the remote index manifest", e); + return Optional.empty(); + } + } + + public void clearLocalIndex() throws IOException { + localIndex.clear(); + } + + public void clearRemoteIndex() throws IOException { + if (remoteIndex != null) { + remoteIndex.clear(); + } + } + + public void notifyIndexState(boolean healthy, IndexNotHealthyException... e) { + if (healthy) { + exception = null; + return; + } + if (e.length > 0 && e[0] != null) { + exception = e[0]; + } + } + + @Override + public void notifyIndexTask(IndexTask task) throws IndexException, IOException { + indexedNid = task.lastEventId; + } + + public void checkIndexState() throws IndexNotHealthyException { + if (exception != null) { + throw exception; + } + } + + public IndexMapping mapping() { + return mapping; + } + + public Directory directory() { + return directory; + } + + public long getIndexedNid() { + return indexedNid; + } + + public void setIndexedNid(long nid) { + this.indexedNid = nid; + } + + @Override + public void close() throws IOException { + directory.close(); + if (remoteIndex instanceof AutoCloseable closeable) { + try { + closeable.close(); + } catch (Exception e) { + if (e instanceof IOException ioe) { + throw ioe; + } + throw new IndexIOException("Failed to close remote index client", e); + } + } + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/IndexSession.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/IndexSession.java new file mode 100644 index 000000000000..4632835f4ec0 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/IndexSession.java @@ -0,0 +1,122 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.index; + +import java.io.IOException; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hive.search.exception.IndexNotReadyException; +import org.apache.hive.search.exception.InitializeException; +import org.apache.hive.search.config.IndexConfig; +import org.apache.hive.search.config.InferenceConfig; +import org.apache.hive.search.config.SearchConfig; +import org.apache.hive.search.inference.EmbedModelRegistry; +import org.apache.hive.search.metastore.MetastoreIndexer; +import org.apache.hive.search.metastore.MetastoreSchemas; +import org.apache.hive.search.metastore.MetastoreTableMapper; +import org.apache.hive.search.search.SearchInternal; +import org.apache.lucene.search.BayesianScoreEstimator; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.SearcherManager; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Colocated indexer and searcher */ +public final class IndexSession implements AutoCloseable { + private static final Logger LOG = LoggerFactory.getLogger(IndexSession.class); + private final Configuration configuration; + private final IndexManager indexManager; + private final EmbedModelRegistry modelRegistry; + private final SearchConfig searchConfig; + + private BayesianScoreEstimator.Parameters parameters; + private MetastoreIndexer metastoreIndexer; + private SearcherManager searcherManager; + private Indexer indexer; + + private final ScheduledExecutorService service = Executors.newScheduledThreadPool(1); + + public IndexSession(Configuration configuration) + throws InitializeException, IOException { + this.configuration = configuration; + IndexConfig indexConfig = new IndexConfig(configuration); + this.searchConfig = new SearchConfig(configuration); + InferenceConfig inferenceConfig = new InferenceConfig(configuration); + this.indexManager = IndexManager.open( + MetastoreSchemas.defaultHiveTablesMapping(indexConfig.indexName(), + inferenceConfig.modelName(), configuration), configuration); + this.modelRegistry = EmbedModelRegistry.create(configuration); + } + + public void maybeRefreshIndex() { + try { + searcherManager.maybeRefresh(); + } catch (IOException e) { + LOG.info("Error while refreshing the index", e); + } + } + + public void initialize() throws Exception { + indexer = new Indexer(indexManager, modelRegistry); + metastoreIndexer = new MetastoreIndexer(configuration, indexManager, indexer); + metastoreIndexer.start(); + searcherManager = new SearcherManager(indexer.writer(), null); + service.scheduleAtFixedRate(this::maybeRefreshIndex, 0, + searchConfig.getRefreshInterval().toSeconds(), TimeUnit.SECONDS); + IndexSearcher searcher = searcherManager.acquire(); + try { + parameters = BayesianScoreEstimator.estimate(searcher, + MetastoreTableMapper.FIELD_SEARCH_TEXT, + searchConfig.getBayesianSamples(), + searchConfig.getBayesianTokensPerQuery(), + searchConfig.getBayesianSeed()); + LOG.info("BayesianScore alpha={} beta={} baseRate={}", + parameters.alpha(), parameters.beta(), parameters.baseRate()); + } finally { + searcherManager.release(searcher); + } + } + + public SearchInternal getSearcher() throws IOException { + if (parameters == null) { + throw new IndexNotReadyException("Index session is not ready for search requests"); + } + indexManager.checkIndexState(); + return new SearchInternal(searcherManager, indexManager, + modelRegistry, searchConfig, parameters); + } + + @Override + public void close() throws Exception { + service.shutdown(); + if (searcherManager != null) { + searcherManager.close(); + } + if (metastoreIndexer != null) { + metastoreIndexer.close(); + } + if (indexer != null) { + indexer.close(); + } + indexManager.close(); + modelRegistry.close(); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/Indexer.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/Indexer.java new file mode 100644 index 000000000000..6ce90b4de0c8 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/Indexer.java @@ -0,0 +1,301 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.index; + +import com.google.common.collect.ArrayListMultimap; +import com.google.common.collect.ListMultimap; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import org.apache.hadoop.hive.common.DatabaseName; +import org.apache.hadoop.hive.metastore.Batchable; +import org.apache.hive.search.config.IndexConfig; +import org.apache.hive.search.exception.IndexException; +import org.apache.hive.search.inference.EmbedModel; +import org.apache.hive.search.inference.EmbedModelRegistry; +import org.apache.hive.search.inference.EmbeddingCache; +import org.apache.hive.search.mapping.FieldSchema; +import org.apache.hive.search.mapping.TableDocument; +import org.apache.hive.search.mapping.field.Field; +import org.apache.hive.search.mapping.field.TextField; +import org.apache.lucene.document.Document; +import org.apache.lucene.index.IndexCommit; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.KeepOnlyLastCommitDeletionPolicy; +import org.apache.lucene.index.SnapshotDeletionPolicy; +import org.apache.lucene.index.Term; +import org.apache.lucene.search.BooleanClause; +import org.apache.lucene.search.BooleanQuery; +import org.apache.lucene.search.PrefixQuery; +import org.apache.lucene.search.Query; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public final class Indexer implements AutoCloseable { + private static final Logger LOG = LoggerFactory.getLogger(Indexer.class); + private static final int EMBED_BATCH_SIZE = 10000; + + private final IndexManager indexManager; + private final EmbedModelRegistry modelRegistry; + private final EmbeddingCache embeddingCache; + private final int commitFlushThreshold; + private SnapshotDeletionPolicy snapshotter; + private FlushTrackingWriter writer; + + public Indexer(IndexManager index, EmbedModelRegistry registry) { + this.indexManager = index; + this.modelRegistry = registry; + this.embeddingCache = registry.embeddingCache(); + this.commitFlushThreshold = + new IndexConfig(index.mapping().configuration()).getCommitFlushes(); + } + + int flushesSinceCommit() { + return writer.flushesSinceCommit(); + } + + public void initialize() throws IOException { + IndexWriterConfig config = + new IndexWriterConfig(indexManager.mapping().analyzer()) + .setCommitOnClose(false) + .setRAMBufferSizeMB(indexManager.mapping().config().getWriteBufferSizeMb()); + SnapshotDeletionPolicy snapshotDeletionPolicy = + new SnapshotDeletionPolicy(new KeepOnlyLastCommitDeletionPolicy()); + config.setIndexDeletionPolicy(snapshotDeletionPolicy); + this.snapshotter = snapshotDeletionPolicy; + this.writer = new FlushTrackingWriter(indexManager.directory(), config); + } + + public void syncBackup() throws IOException { + if (!indexManager.hasBackup()) { + return; + } + IndexCommit snapshot = snapshotter.snapshot(); + try { + indexManager.syncBackup(); + } finally { + snapshotter.release(snapshot); + } + } + + /** Writes already-embedded documents to Lucene. */ + private void writeDocuments(List docs) throws IOException, IndexException { + List luceneDocs = new ArrayList<>(); + List ids = new ArrayList<>(); + for (TableDocument doc : docs) { + ids.add(doc.idField().value()); + luceneDocs.addAll(doc.toDocuments()); + } + delete(ids.toArray(new String[0])); + writer.addDocuments(luceneDocs); + } + + public List embedDocuments(List tableDocs) + throws IndexException { + long start = System.currentTimeMillis(); + List result = new ArrayList<>(tableDocs.size()); + Map> modelPerTxt = new HashMap<>(); + for (TableDocument doc : tableDocs) { + TableDocument newDoc = new TableDocument(doc.idField(), List.of(), indexManager.mapping()); + result.add(newDoc); + for (Field field : doc.fields()) { + if (field instanceof org.apache.hive.search.mapping.field.IdField) { + continue; + } + if (field instanceof TextField text) { + FieldSchema schema = indexManager.mapping().fieldSchema(text.name()); + if (schema instanceof FieldSchema.TextFieldSchema textSchema + && textSchema.search().semantic()) { + if (text.embedding() != null) { + newDoc.appendField(field); + } else { + String modelRef = textSchema.search().semanticModel(); + ListMultimap multimapText = + modelPerTxt.computeIfAbsent(modelRef, s -> ArrayListMultimap.create()); + multimapText.put(text, newDoc); + } + } else { + newDoc.appendField(field); + } + } else { + newDoc.appendField(field); + } + } + } + int totalFields = modelPerTxt.values().stream().mapToInt(ListMultimap::size).sum(); + long cacheHitsBefore = embeddingCache.hits(); + long cacheMissesBefore = embeddingCache.misses(); + for (Map.Entry> entry : modelPerTxt.entrySet()) { + embedInBatch(entry.getKey(), modelRegistry.get(entry.getKey()), entry.getValue()); + } + if (totalFields > 0) { + long cacheHits = embeddingCache.hits() - cacheHitsBefore; + long cacheMisses = embeddingCache.misses() - cacheMissesBefore; + LOG.info("Embedded {} semantic field(s) across {} document(s) in {}ms" + + " (embedding cache hits={}, misses={})", + totalFields, tableDocs.size(), System.currentTimeMillis() - start, cacheHits, cacheMisses); + } + return result; + } + + private void embedInBatch(String modelRef, EmbedModel embedModel, + ListMultimap textDocs) throws IndexException { + int uniqueTexts = textDocs.keySet().stream() + .map(TextField::value) + .collect(java.util.stream.Collectors.toSet()) + .size(); + long start = System.currentTimeMillis(); + try { + Batchable.runBatched(EMBED_BATCH_SIZE, new ArrayList<>(textDocs.keySet()), + new Batchable() { + @Override + public List run(List batchFields) throws IndexException { + long batchStart = System.currentTimeMillis(); + ListMultimap valueToTxt = ArrayListMultimap.create(); + batchFields.forEach(f -> valueToTxt.put(f.value(), f)); + List missTexts = new ArrayList<>(); + for (String text : valueToTxt.keySet()) { + Optional cached = + embeddingCache.get(modelRef, EmbedModel.TaskType.DOCUMENT, text); + if (cached.isPresent()) { + applyEmbedding(valueToTxt, textDocs, text, cached.get()); + } else { + missTexts.add(text); + } + } + if (!missTexts.isEmpty()) { + String[] texts = missTexts.toArray(new String[0]); + float[][] embeddings = embedModel.embedBatch(EmbedModel.TaskType.DOCUMENT, texts); + for (int i = 0; i < texts.length; i++) { + String text = texts[i]; + float[] embedding = embeddings[i]; + embeddingCache.put(modelRef, EmbedModel.TaskType.DOCUMENT, text, embedding); + applyEmbedding(valueToTxt, textDocs, text, embedding); + } + } + LOG.debug("Model '{}' embedded batch of {} unique text(s), {} cache miss(es) in {}ms", + modelRef, valueToTxt.keySet().size(), missTexts.size(), + System.currentTimeMillis() - batchStart); + return List.of(); + } + }); + } catch (Exception e) { + throw IndexException.wrap("Error while embedding the documents with model '" + modelRef + "'", + e); + } + LOG.info("Model '{}' embedded {} field(s) from {} unique text(s) in {}ms", + modelRef, textDocs.size(), uniqueTexts, System.currentTimeMillis() - start); + } + + private void applyEmbedding(ListMultimap valueToTxt, + ListMultimap textDocs, String text, float[] embedding) { + for (TextField tf : valueToTxt.get(text)) { + for (TableDocument document : textDocs.get(tf)) { + document.appendField(tf.withEmbedding(embedding)); + } + } + } + + public void addDocuments(List docs) throws IOException, IndexException { + writeDocuments(embedDocuments(docs)); + } + + public IndexWriter writer() { + return writer; + } + + public EmbeddingCache embeddingCache() { + return embeddingCache; + } + + /** + * Commits the index when forced or after enough Lucene auto-flushes. + * Lucene continues to flush segments automatically based on RAM buffer settings. + */ + public boolean flush(long lastEventId, boolean force) + throws IOException { + if (!force && (!hasPendingChanges() || !shouldCommit())) { + return false; + } + String model = indexManager.mapping().inference().modelName(); + Map metadata = Map.of( + "nid", lastEventId + "", + "model", model, + "commit_time", String.valueOf(System.currentTimeMillis()) + ); + writer.setLiveCommitData(metadata.entrySet()); + long seqnum = writer.commit(); + if (seqnum < 0) { + return false; + } + writer.resetFlushTracking(); + return true; + } + + private boolean shouldCommit() { + return commitFlushThreshold <= 0 || writer.flushesSinceCommit() >= commitFlushThreshold; + } + + private boolean hasPendingChanges() { + return writer.hasUncommittedChanges() + || writer.numRamDocs() > 0 + || writer.hasDeletions(); + } + + public int delete(String... docIds) throws IOException { + if (docIds == null || docIds.length == 0) { + return 0; + } + int before = writer.getDocStats().numDocs; + Term[] terms = Arrays.stream(docIds) + .map(id -> new Term("_id" + TableDocument.FILTER_SUFFIX, id)) + .toArray(Term[]::new); + writer.deleteDocuments(terms); + int after = writer.getDocStats().numDocs; + return before - after; + } + + public int deleteDatabases(DatabaseName... databases) + throws IOException { + if (databases == null || databases.length == 0) { + return 0; + } + int before = writer.getDocStats().numDocs; + BooleanQuery.Builder builder = new BooleanQuery.Builder(); + for (DatabaseName database : databases) { + Query query = new PrefixQuery(new Term("_id" + TableDocument.FILTER_SUFFIX, + database.getCat() + "." + database.getDb() + ".")); + builder.add(query, BooleanClause.Occur.SHOULD); + } + writer.deleteDocuments(builder.build()); + int after = writer.getDocStats().numDocs; + return before - after; + } + + @Override + public void close() throws IOException { + writer.close(); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/manifest/IndexManifest.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/manifest/IndexManifest.java new file mode 100644 index 000000000000..9f9544003620 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/manifest/IndexManifest.java @@ -0,0 +1,100 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.index.manifest; + +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.databind.ObjectMapper; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; + +@JsonIgnoreProperties(ignoreUnknown = true) +public record IndexManifest(String indexName, List files, String modelName, long lastEventId) { + public static final String MANIFEST_FILE_NAME = "index.json"; + /** Local-only staging copy of the remote manifest while a restore is in progress. */ + public static final String STAGING_MANIFEST_FILE_NAME = ".index.json"; + private static final ObjectMapper JSON = new ObjectMapper(); + + public record IndexFile(String name, long size) {} + + public static IndexManifest create(String indexName, List files, + String modelName, long lastEventId) { + return new IndexManifest(indexName, List.copyOf(files), modelName, lastEventId); + } + + public byte[] toJsonBytes() throws IOException { + return JSON.writerWithDefaultPrettyPrinter().writeValueAsBytes(this); + } + + public static IndexManifest fromJson(byte[] bytes) throws IOException { + return JSON.readValue(bytes, IndexManifest.class); + } + + public boolean sameFilesAs(IndexManifest other) { + return other != null && toMap(files).equals(toMap(other.files())); + } + + public List diff(IndexManifest target) { + Map sourceMap = toMap(files); + Map destMap = target == null ? Map.of() : toMap(target.files()); + HashSet keys = new HashSet<>(); + keys.addAll(sourceMap.keySet()); + keys.addAll(destMap.keySet()); + + List ops = new ArrayList<>(); + for (String key : keys) { + Long sourceSize = sourceMap.get(key); + Long destSize = destMap.get(key); + if (sourceSize != null && sourceSize.equals(destSize)) { + continue; + } + if (sourceSize != null) { + ops.add(ChangedFileOp.add(key, sourceSize)); + } else if (destSize != null) { + ops.add(ChangedFileOp.del(key)); + } + } + return ops; + } + + private static Map toMap(List files) { + HashMap map = new HashMap<>(); + for (IndexFile file : files) { + map.put(file.name(), file.size()); + } + return map; + } + + public interface ChangedFileOp { + record Add(String fileName, Long size) implements ChangedFileOp {} + + record Del(String fileName) implements ChangedFileOp {} + + static Add add(String fileName, Long size) { + return new Add(fileName, size); + } + + static Del del(String fileName) { + return new Del(fileName); + } + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/store/IndexBackupUtils.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/store/IndexBackupUtils.java new file mode 100644 index 000000000000..a467d6e8135d --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/store/IndexBackupUtils.java @@ -0,0 +1,121 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.index.store; + +import java.io.IOException; +import java.io.InputStream; +import java.net.URI; +import java.util.List; +import java.util.Optional; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hive.search.config.IndexStoreConfig; +import org.apache.hive.search.exception.IndexIOException; +import org.apache.hive.search.index.manifest.IndexManifest; + +public final class IndexBackupUtils { + private IndexBackupUtils() {} + + /** Push local index state to the backup (remote) location. */ + public static boolean syncToBackup(IndexStateClient local, IndexStateClient remote) + throws IOException { + Optional localManifest = local.readManifest(); + if (localManifest.isEmpty()) { + return false; + } + Optional remoteManifest = remote.readManifest(); + if (remoteManifest.isPresent() + && remoteManifest.get().lastEventId() >= localManifest.get().lastEventId()) { + return false; + } + applyDiff(local, remote, localManifest.get().diff(remoteManifest.orElse(null))); + return remote.writeManifest(localManifest.get()); + } + + /** + * If a previous restore left a staging manifest, finalize it when files already match, + * or keep it in place so {@link #restoreFromBackup} can resume. + */ + public static void resolveInterruptedRestore(IndexStateClient local) + throws IOException { + Optional staging = local.readStagingManifest(); + if (staging.isEmpty()) { + local.clearStagingManifest(); + return; + } + IndexManifest target = staging.get(); + if (target.sameFilesAs(local.readLocalFileManifest())) { + local.validateRestoredIndex(target); + local.clearStagingManifest(); + } + } + + /** Pull backup (remote) index state into the local directory. */ + public static boolean restoreFromBackup(IndexStateClient local, IndexStateClient remote) + throws IOException { + Optional remoteManifest = remote.readManifest(); + if (remoteManifest.isEmpty()) { + return false; + } + IndexManifest target = remoteManifest.get(); + Optional localManifest = local.readManifest(); + if (localManifest.isPresent() + && localManifest.get().lastEventId() >= target.lastEventId()) { + return false; + } + prepareStagingManifest(local, target); + IndexManifest localFiles = local.readLocalFileManifest(); + applyDiff(remote, local, target.diff(localFiles)); + local.validateRestoredIndex(target); + local.clearStagingManifest(); + return true; + } + + private static void prepareStagingManifest(IndexStateClient local, IndexManifest target) + throws IOException { + Optional staging = local.readStagingManifest(); + if (staging.isPresent() + && staging.get().lastEventId() == target.lastEventId() + && staging.get().modelName().equals(target.modelName())) { + return; + } + local.clearStagingManifest(); + local.writeStagingManifest(target); + } + + private static void applyDiff(IndexStateClient source, IndexStateClient dest, + List ops) throws IOException { + for (IndexManifest.ChangedFileOp op : ops) { + switch (op) { + case IndexManifest.ChangedFileOp.Add add -> { + try (InputStream in = source.read(add.fileName())) { + dest.write(add.fileName(), in); + } + } + case IndexManifest.ChangedFileOp.Del del -> dest.delete(del.fileName()); + default -> throw new IndexIOException("Unexpected file operation during backup sync: " + op); + } + } + } + + public static IndexStateClient openRemote(String remoteUri, String indexName, Configuration conf) + throws IOException { + IndexStoreConfig.validateRemoteUri(remoteUri); + return new RemoteStateClient(URI.create(remoteUri), conf, indexName); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/store/IndexStateClient.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/store/IndexStateClient.java new file mode 100644 index 000000000000..bf69dda471bb --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/store/IndexStateClient.java @@ -0,0 +1,65 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.index.store; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.util.Optional; + +import org.apache.hive.search.exception.IndexIOException; +import org.apache.hive.search.index.manifest.IndexManifest; + +public interface IndexStateClient { + + Optional readManifest() throws IOException; + + InputStream read(String fileName) throws IOException; + + void write(String fileName, InputStream stream) throws IOException; + + void delete(String fileName) throws IOException; + + default boolean writeManifest(IndexManifest manifest) throws IOException { + write(IndexManifest.MANIFEST_FILE_NAME, new ByteArrayInputStream(manifest.toJsonBytes())); + return true; + } + + default Optional readStagingManifest() throws IOException { + return Optional.empty(); + } + + default void writeStagingManifest(IndexManifest manifest) throws IOException { + + } + + default void clearStagingManifest() throws IOException { + + } + + /** Snapshot of index files currently on disk, excluding staging metadata */ + default IndexManifest readLocalFileManifest() throws IOException { + return readManifest().orElse(null); + } + + default void validateRestoredIndex(IndexManifest expected) throws IndexIOException { + + } + + void clear() throws IOException; +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/store/LocalStateClient.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/store/LocalStateClient.java new file mode 100644 index 000000000000..15801da921e2 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/store/LocalStateClient.java @@ -0,0 +1,259 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.index.store; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.nio.file.NoSuchFileException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import org.apache.commons.lang3.StringUtils; +import org.apache.hive.search.exception.IndexIOException; +import org.apache.hive.search.index.manifest.IndexManifest; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.SegmentInfos; +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.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public record LocalStateClient(Directory directory, String indexName) + implements IndexStateClient { + private static final Logger LOG = LoggerFactory.getLogger(LocalStateClient.class); + + public static boolean isStagingFile(String fileName) { + return IndexManifest.STAGING_MANIFEST_FILE_NAME.equals(fileName); + } + + @Override + public Optional readManifest() throws IOException { + if (hasStagingManifest()) { + return Optional.empty(); + } + if (!DirectoryReader.indexExists(directory)) { + return Optional.empty(); + } + SegmentInfos segmentInfos = SegmentInfos.readLatestCommit(directory); + Map userData = segmentInfos.getUserData(); + String nid = userData.get("nid"); + if (StringUtils.isEmpty(nid)) { + return Optional.empty(); + } + String modelName = userData.get("model"); + if (StringUtils.isEmpty(modelName)) { + return Optional.empty(); + } + long eventId; + try { + eventId = Long.parseLong(nid); + } catch (NumberFormatException e) { + throw new IndexIOException("Invalid commit checkpoint in index metadata: nid=" + nid, e); + } + List files = new ArrayList<>(); + for (String file : segmentInfos.files(true)) { + files.add(new IndexManifest.IndexFile(file, fileLength(file))); + } + return Optional.of(IndexManifest.create(indexName, files, modelName, eventId)); + } + + @Override + public Optional readStagingManifest() throws IOException { + if (!hasStagingManifest()) { + return Optional.empty(); + } + try (IndexInput in = directory.openInput(IndexManifest.STAGING_MANIFEST_FILE_NAME, + IOContext.DEFAULT)) { + byte[] bytes = new byte[(int) in.length()]; + in.readBytes(bytes, 0, bytes.length); + return Optional.of(IndexManifest.fromJson(bytes)); + } catch (IOException e) { + LOG.warn("Failed to read the staging manifest", e); + return Optional.empty(); + } + } + + @Override + public void writeStagingManifest(IndexManifest manifest) throws IOException { + write(IndexManifest.STAGING_MANIFEST_FILE_NAME, + new ByteArrayInputStream(manifest.toJsonBytes())); + } + + @Override + public void clearStagingManifest() throws IOException { + delete(IndexManifest.STAGING_MANIFEST_FILE_NAME); + } + + public boolean hasStagingManifest() throws IOException { + return Arrays.asList(directory.listAll()) + .contains(IndexManifest.STAGING_MANIFEST_FILE_NAME); + } + + @Override + public IndexManifest readLocalFileManifest() throws IOException { + List files = new ArrayList<>(); + for (String file : directory.listAll()) { + if (isStagingFile(file)) { + continue; + } + files.add(new IndexManifest.IndexFile(file, fileLength(file))); + } + return IndexManifest.create(indexName, files, "", -1); + } + + public boolean isIndexReadable() throws IOException { + return DirectoryReader.indexExists(directory); + } + + @Override + public void validateRestoredIndex(IndexManifest expected) throws IndexIOException { + try { + if (!expected.sameFilesAs(readLocalFileManifest())) { + throw new IndexIOException("Restored index files do not match staging manifest"); + } + if (!isIndexReadable()) { + throw new IndexIOException("Restored index is not readable by Lucene"); + } + SegmentInfos segmentInfos = SegmentInfos.readLatestCommit(directory); + Map userData = segmentInfos.getUserData(); + String nid = userData.get("nid"); + if (StringUtils.isEmpty(nid)) { + throw new IndexIOException("Restored index is missing commit checkpoint"); + } + String modelName = userData.get("model"); + if (StringUtils.isEmpty(modelName)) { + throw new IndexIOException("Restored index is missing embedding model metadata"); + } + long eventId; + try { + eventId = Long.parseLong(nid); + } catch (NumberFormatException e) { + throw new IndexIOException("Restored index has invalid commit checkpoint: nid=" + nid, e); + } + if (eventId != expected.lastEventId()) { + throw new IndexIOException( + "Restored index checkpoint mismatch: expected nid=" + expected.lastEventId() + + " but found nid=" + eventId); + } + if (!expected.modelName().equals(modelName)) { + throw new IndexIOException("Restored index embedding model mismatch"); + } + } catch (IOException e) { + throw IndexIOException.wrap(e); + } + } + + @Override + public InputStream read(String fileName) throws IOException { + return new DirectoryInputStream(directory, fileName); + } + + @Override + public void write(String fileName, InputStream stream) throws IOException { + if (Arrays.asList(directory.listAll()).contains(fileName)) { + directory.deleteFile(fileName); + } + try (IndexOutput out = directory.createOutput(fileName, IOContext.DEFAULT)) { + stream.transferTo(new OutputStreamAdapter(out)); + } + } + + @Override + public void delete(String fileName) throws IOException { + if (Arrays.asList(directory.listAll()).contains(fileName)) { + directory.deleteFile(fileName); + } + } + + @Override + public void clear() throws IOException { + for (String file : directory.listAll()) { + directory.deleteFile(file); + } + } + + private long fileLength(String fileName) throws IOException { + try (IndexInput input = directory.openInput(fileName, IOContext.DEFAULT)) { + return input.length(); + } + } + + private static final class DirectoryInputStream extends InputStream { + private final IndexInput input; + + DirectoryInputStream(Directory directory, String fileName) throws IOException { + if (!Arrays.asList(directory.listAll()).contains(fileName)) { + throw new NoSuchFileException(fileName); + } + this.input = directory.openInput(fileName, IOContext.DEFAULT); + } + + @Override + public int read() throws IOException { + if (input.getFilePointer() >= input.length()) { + return -1; + } + return Byte.toUnsignedInt(input.readByte()); + } + + @Override + public int read(byte[] b, int off, int len) throws IOException { + int remaining = (int) Math.min(len, input.length() - input.getFilePointer()); + if (remaining <= 0) { + return -1; + } + input.readBytes(b, off, remaining); + return remaining; + } + + @Override + public void close() throws IOException { + input.close(); + } + } + + private static final class OutputStreamAdapter extends OutputStream { + private final IndexOutput output; + + OutputStreamAdapter(IndexOutput output) { + this.output = output; + } + + @Override + public void write(int b) throws IOException { + output.writeByte((byte) b); + } + + @Override + public void write(byte[] b, int off, int len) throws IOException { + output.writeBytes(b, off, len); + } + + @Override + public void close() throws IOException { + output.close(); + } + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/store/RemoteStateClient.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/store/RemoteStateClient.java new file mode 100644 index 000000000000..4067d7e4c5f4 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/store/RemoteStateClient.java @@ -0,0 +1,88 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.index.store; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.net.URI; +import java.util.Optional; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.FileSystem; +import org.apache.hadoop.fs.Path; +import org.apache.hive.search.index.manifest.IndexManifest; + +/** Remote index backup target backed by Hadoop {@link FileSystem}. */ +public final class RemoteStateClient implements IndexStateClient { + private final FileSystem fs; + private final Path root; + + public RemoteStateClient(URI baseUri, Configuration conf, String indexName) + throws IOException { + this.fs = FileSystem.get(baseUri, conf); + this.root = resolveRoot(baseUri, indexName); + fs.mkdirs(root); + } + + private static Path resolveRoot(URI baseUri, String indexName) { + return new Path(new Path(baseUri), indexName); + } + + @Override + public Optional readManifest() throws IOException { + Path manifestPath = new Path(root, IndexManifest.MANIFEST_FILE_NAME); + if (!fs.exists(manifestPath)) { + return Optional.empty(); + } + try (InputStream in = fs.open(manifestPath)) { + return Optional.of(IndexManifest.fromJson(in.readAllBytes())); + } + } + + @Override + public InputStream read(String fileName) throws IOException { + return fs.open(new Path(root, fileName)); + } + + @Override + public void write(String fileName, InputStream stream) throws IOException { + Path target = new Path(root, fileName); + Path parent = target.getParent(); + if (parent != null && !fs.exists(parent)) { + fs.mkdirs(parent); + } + try (OutputStream out = fs.create(target, true)) { + stream.transferTo(out); + } + } + + @Override + public void delete(String fileName) throws IOException { + Path target = new Path(root, fileName); + if (fs.exists(target)) { + fs.delete(target, false); + } + } + + @Override + public void clear() throws IOException { + fs.delete(root, true); + fs.mkdirs(root); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/EmbedModel.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/EmbedModel.java new file mode 100644 index 000000000000..33ebf31a303a --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/EmbedModel.java @@ -0,0 +1,39 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.inference; + +import org.apache.hive.search.exception.IndexException; + +public interface EmbedModel extends AutoCloseable { + enum TaskType { + DOCUMENT, + QUERY + } + + String name(); + + float[] embed(TaskType task, String text) throws IndexException; + + default float[][] embedBatch(TaskType task, String[] texts) throws IndexException { + float[][] result = new float[texts.length][]; + for (int i = 0; i < texts.length; i++) { + result[i] = embed(task, texts[i]); + } + return result; + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/EmbedModelRegistry.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/EmbedModelRegistry.java new file mode 100644 index 000000000000..268070f1fac7 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/EmbedModelRegistry.java @@ -0,0 +1,76 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.inference; + +import java.io.IOException; +import java.util.Map; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hive.search.config.InferenceConfig; +import org.apache.hive.search.exception.IndexException; +import org.apache.hive.search.exception.InitializeException; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public record EmbedModelRegistry(Map models, EmbeddingCache embeddingCache) + implements AutoCloseable { + private static final Logger LOG = LoggerFactory.getLogger(EmbedModelRegistry.class); + + public EmbedModelRegistry(Map models) { + this(models, EmbeddingCache.disabled()); + } + + public EmbedModelRegistry(Map models, EmbeddingCache embeddingCache) { + this.models = Map.copyOf(models); + this.embeddingCache = embeddingCache; + } + + public static EmbedModelRegistry create(Configuration configuration) + throws InitializeException, IOException { + InferenceConfig inference = new InferenceConfig(configuration); + EmbeddingCache embeddingCache = EmbeddingCache.create(configuration); + long start = System.currentTimeMillis(); + String modelName = inference.modelName(); + EmbedModel embedModel = new LocalOnnxEmbeddingModel(inference); + long warmupStart = System.currentTimeMillis(); + try { + embedModel.embed(EmbedModel.TaskType.QUERY, "warmup"); + } catch (IndexException e) { + throw new InitializeException("Failed to warm up embedding model '" + modelName + "'", e); + } + LOG.info("Loaded embedding model '{}' in {}ms (warmup {}ms)", + modelName, System.currentTimeMillis() - start, + System.currentTimeMillis() - warmupStart); + return new EmbedModelRegistry(Map.of(modelName, embedModel), embeddingCache); + } + + public EmbedModel get(String modelRef) throws IndexException { + EmbedModel model = models.get(modelRef); + if (model == null) { + throw new IndexException("Embedding model '" + modelRef + "' is not configured"); + } + return model; + } + + @Override + public void close() throws Exception { + for (EmbedModel model : models.values()) { + model.close(); + } + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/EmbedRequest.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/EmbedRequest.java new file mode 100644 index 000000000000..8cbfd31c2aae --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/EmbedRequest.java @@ -0,0 +1,80 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.inference; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.atomic.AtomicReference; + +import org.apache.hive.search.exception.IndexException; + +import ai.onnxruntime.OrtException; + +/** One embedding job waiting on the shared inference queue. */ +final class EmbedRequest { + static final EmbedRequest SHUTDOWN = new EmbedRequest(null, null); + + private final EmbedModel.TaskType task; + private final String text; + private final CountDownLatch done = new CountDownLatch(1); + private final AtomicReference result = new AtomicReference<>(); + private final AtomicReference error = new AtomicReference<>(); + + EmbedRequest(EmbedModel.TaskType task, String text) { + this.task = task; + this.text = text; + } + + EmbedModel.TaskType task() { + return task; + } + + String text() { + return text; + } + + boolean isShutdown() { + return this == SHUTDOWN; + } + + void complete(float[] vector) { + result.set(vector); + done.countDown(); + } + + void fail(Exception failure) { + error.set(failure); + done.countDown(); + } + + float[] awaitResult() throws IndexException { + try { + done.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw IndexException.wrap("Embedding interrupted", e); + } + Exception failure = error.get(); + if (failure != null) { + if (failure instanceof OrtException ort) { + throw IndexException.wrap("Failed to encode text", ort); + } + throw IndexException.wrap("Failed to encode text", failure); + } + return result.get(); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/EmbeddingCache.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/EmbeddingCache.java new file mode 100644 index 000000000000..e3f3757c9693 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/EmbeddingCache.java @@ -0,0 +1,99 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.inference; + +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicLong; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hive.search.config.InferenceConfig; + +import com.google.common.cache.Cache; +import com.google.common.cache.CacheBuilder; +import com.google.common.hash.Hashing; + +/** In-memory LRU cache of document embeddings keyed by model, task, and text hash. */ +public final class EmbeddingCache { + private final Cache cache; + private final AtomicLong hits = new AtomicLong(); + private final AtomicLong misses = new AtomicLong(); + + private EmbeddingCache(Cache cache) { + this.cache = cache; + } + + public static EmbeddingCache create(Configuration configuration) { + InferenceConfig inference = new InferenceConfig(configuration); + if (!inference.isEmbeddingCacheEnabled()) { + return disabled(); + } + return new EmbeddingCache(CacheBuilder.newBuilder() + .maximumSize(inference.getEmbeddingCacheMaxEntries()) + .build()); + } + + public static EmbeddingCache disabled() { + return new EmbeddingCache(null); + } + + public Optional get(String modelName, EmbedModel.TaskType task, String text) { + if (cache == null) { + misses.incrementAndGet(); + return Optional.empty(); + } + CacheKey key = CacheKey.of(modelName, task, text); + float[] cached = cache.getIfPresent(key); + if (cached == null) { + misses.incrementAndGet(); + return Optional.empty(); + } + hits.incrementAndGet(); + return Optional.of(Arrays.copyOf(cached, cached.length)); + } + + public void put(String modelName, EmbedModel.TaskType task, String text, float[] embedding) { + if (cache == null) { + return; + } + cache.put(CacheKey.of(modelName, task, text), Arrays.copyOf(embedding, embedding.length)); + } + + public long hits() { + return hits.get(); + } + + public long misses() { + return misses.get(); + } + + public boolean enabled() { + return cache != null; + } + + private record CacheKey(String modelName, EmbedModel.TaskType task, long textHash, + String text) { + static CacheKey of(String modelName, EmbedModel.TaskType task, String text) { + long hash = Hashing.murmur3_128() + .hashString(text, StandardCharsets.UTF_8) + .asLong(); + return new CacheKey(modelName, task, hash, text); + } + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/EmbeddingPrompt.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/EmbeddingPrompt.java new file mode 100644 index 000000000000..41f5689aa2b2 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/EmbeddingPrompt.java @@ -0,0 +1,36 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.inference; + +/** Model-specific text prefixes applied before tokenization. */ +public record EmbeddingPrompt(String documentPrefix, String queryPrefix) { + public static EmbeddingPrompt e5() { + return new EmbeddingPrompt("passage: ", "query: "); + } + + public static EmbeddingPrompt none() { + return new EmbeddingPrompt("", ""); + } + + public String prefixFor(EmbedModel.TaskType task) { + return switch (task) { + case DOCUMENT -> documentPrefix; + case QUERY -> queryPrefix; + }; + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/InferenceWorker.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/InferenceWorker.java new file mode 100644 index 000000000000..6a10baf9d2a0 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/InferenceWorker.java @@ -0,0 +1,242 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.inference; + +import ai.djl.huggingface.tokenizers.Encoding; +import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer; +import ai.onnxruntime.OnnxTensor; +import ai.onnxruntime.OrtEnvironment; +import ai.onnxruntime.OrtException; +import ai.onnxruntime.OrtSession; + +import java.io.IOException; +import java.nio.LongBuffer; +import java.nio.file.Path; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.TimeUnit; + +import org.apache.hive.search.config.InferenceConfig; +import org.apache.hive.search.exception.InitializeException; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Worker thread that pulls embedding jobs from a shared queue. */ +final class InferenceWorker extends Thread { + private static final Logger LOG = LoggerFactory.getLogger(InferenceWorker.class); + + private final OrtEnvironment ortEnv; + private final OrtSession session; + private final HuggingFaceTokenizer tokenizer; + private final EmbeddingPrompt prompt; + private final BlockingQueue queue; + private final int maxSeqLength; + + InferenceWorker( + InferenceConfig config, + InferenceConfig.EmbeddingModelSpec spec, + int workerIndex, + BlockingQueue queue) + throws InitializeException, IOException { + super("EmbedModel-" + spec.getModel() + "-" + workerIndex); + setDaemon(true); + this.prompt = spec.getPrompt() == null ? EmbeddingPrompt.none() : spec.getPrompt(); + this.queue = queue; + this.maxSeqLength = config.getEmbeddingMaxSeqLength(); + int intraOpThreads = Math.max(1, + Runtime.getRuntime().availableProcessors() / config.getEmbeddingThreads()); + Path modelDir = spec.getModelDir(); + try { + this.ortEnv = OrtEnvironment.getEnvironment(); + OrtSession.SessionOptions opts = new OrtSession.SessionOptions(); + opts.setIntraOpNumThreads(intraOpThreads); + this.session = ortEnv.createSession( + modelDir.resolve(InferenceConfig.MODEL_ONNX_FILE).toString(), opts); + this.tokenizer = HuggingFaceTokenizer.newInstance( + modelDir.resolve(InferenceConfig.TOKENIZER)); + } catch (OrtException e) { + throw InitializeException.wrap("Failed to initialize embedding worker", e); + } + } + + void startWorker() { + start(); + } + + void awaitStop() { + try { + join(TimeUnit.SECONDS.toMillis(30)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + interrupt(); + } + } + + @Override + public void run() { + try { + while (true) { + EmbedRequest request; + try { + request = queue.take(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + break; + } + if (request.isShutdown()) { + break; + } + try { + request.complete(embedInternal(request.task(), request.text())); + } catch (Exception e) { + request.fail(e); + } + } + } finally { + closeResources(); + } + } + + private float[] embedInternal(EmbedModel.TaskType task, String text) throws OrtException { + String input = prompt.prefixFor(task) + text; + Encoding encoding = tokenizer.encode(input, true, true); + long[] inputIds = encoding.getIds(); + long[] attentionMask = encoding.getAttentionMask(); + + if (inputIds.length <= maxSeqLength) { + return normalize(embedInternal(inputIds, attentionMask, 0, inputIds.length)); + } + + int chunkCount = (inputIds.length + maxSeqLength - 1) / maxSeqLength; + float[][] chunkVectors = new float[chunkCount][]; + for (int chunk = 0; chunk < chunkCount; chunk++) { + int offset = chunk * maxSeqLength; + int chunkLen = Math.min(maxSeqLength, inputIds.length - offset); + chunkVectors[chunk] = normalize(embedInternal(inputIds, attentionMask, offset, chunkLen)); + } + if (LOG.isDebugEnabled()) { + LOG.debug("Embedded {} token(s) in {} chunk(s) on worker {}", inputIds.length, chunkCount, + getName()); + } + return normalize(meanPoolVectors(chunkVectors)); + } + + private float[] embedInternal(long[] inputIds, long[] attentionMask, int offset, int seqLen) + throws OrtException { + long[] ids = new long[seqLen]; + long[] mask = new long[seqLen]; + long[] tokenTypeIds = new long[seqLen]; + System.arraycopy(inputIds, offset, ids, 0, seqLen); + System.arraycopy(attentionMask, offset, mask, 0, seqLen); + + long[] shape = {1, seqLen}; + try (OnnxTensor idsTensor = OnnxTensor.createTensor(ortEnv, LongBuffer.wrap(ids), shape); + OnnxTensor maskTensor = OnnxTensor.createTensor(ortEnv, LongBuffer.wrap(mask), shape); + OnnxTensor typeTensor = OnnxTensor.createTensor(ortEnv, LongBuffer.wrap(tokenTypeIds), shape)) { + + Map inputs = new HashMap<>(); + inputs.put("input_ids", idsTensor); + inputs.put("attention_mask", maskTensor); + inputs.put("token_type_ids", typeTensor); + try (OrtSession.Result result = session.run(inputs)) { + return poolOutput(result.get(0).getValue(), mask, seqLen); + } + } + } + + private static float[] poolOutput(Object value, long[] mask, int seqLen) { + if (value instanceof float[][][] tokenEmbeddings) { + return meanPool(tokenEmbeddings[0], mask, seqLen); + } + if (value instanceof float[][] matrix) { + if (matrix.length == 1) { + return matrix[0].clone(); + } + return meanPool(matrix, mask, seqLen); + } + if (value instanceof float[] vector) { + return vector.clone(); + } + throw new IllegalStateException("Unsupported ONNX embedding output type: " + value.getClass()); + } + + static float[] meanPool(float[][] tokenEmbeddings, long[] mask, int seqLen) { + int dim = tokenEmbeddings[0].length; + float[] result = new float[dim]; + float maskSum = 0; + for (int i = 0; i < seqLen && i < tokenEmbeddings.length; i++) { + if (mask[i] == 1) { + maskSum++; + for (int j = 0; j < dim; j++) { + result[j] += tokenEmbeddings[i][j]; + } + } + } + if (maskSum > 0) { + for (int j = 0; j < dim; j++) { + result[j] /= maskSum; + } + } + return result; + } + + static float[] meanPoolVectors(float[][] vectors) { + if (vectors.length == 1) { + return vectors[0].clone(); + } + int dim = vectors[0].length; + float[] result = new float[dim]; + for (float[] vector : vectors) { + for (int j = 0; j < dim; j++) { + result[j] += vector[j]; + } + } + for (int j = 0; j < dim; j++) { + result[j] /= vectors.length; + } + return result; + } + + static float[] normalize(float[] vec) { + float norm = 0; + for (float v : vec) { + norm += v * v; + } + norm = (float) Math.sqrt(norm); + if (norm > 0) { + for (int i = 0; i < vec.length; i++) { + vec[i] /= norm; + } + } + return vec; + } + + private void closeResources() { + if (tokenizer != null) { + tokenizer.close(); + } + if (session != null) { + try { + session.close(); + } catch (OrtException e) { + LOG.warn("Failed to close ONNX session on worker {}", getName(), e); + } + } + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/LocalOnnxEmbeddingModel.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/LocalOnnxEmbeddingModel.java new file mode 100644 index 000000000000..a8e59c6476a3 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/LocalOnnxEmbeddingModel.java @@ -0,0 +1,105 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.inference; + +import java.io.IOException; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.LinkedBlockingQueue; + +import org.apache.hive.search.config.InferenceConfig; +import org.apache.hive.search.exception.IndexException; +import org.apache.hive.search.exception.InitializeException; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** ONNX Runtime-backed embedding model for locally deployed sentence-transformer models. */ +public final class LocalOnnxEmbeddingModel implements EmbedModel { + private static final Logger LOG = LoggerFactory.getLogger(LocalOnnxEmbeddingModel.class); + + private final String name; + private final InferenceWorker[] workers; + private final BlockingQueue queue = new LinkedBlockingQueue<>(); + + public LocalOnnxEmbeddingModel(InferenceConfig config) throws InitializeException, IOException { + InferenceConfig.EmbeddingModelSpec spec = config.embedding(); + this.name = spec.getModel(); + int threads = config.getEmbeddingThreads(); + this.workers = new InferenceWorker[threads]; + for (int i = 0; i < threads; i++) { + workers[i] = new InferenceWorker(config, spec, i, queue); + workers[i].startWorker(); + } + LOG.info("Loaded ONNX embedding model '{}' from {} with {} worker thread(s), max {} token(s) per chunk", + name, spec.getModelDir(), threads, config.getEmbeddingMaxSeqLength()); + } + + @Override + public String name() { + return name; + } + + @Override + public float[] embed(TaskType task, String text) throws IndexException { + return enqueue(task, text).awaitResult(); + } + + @Override + public float[][] embedBatch(TaskType task, String[] texts) throws IndexException { + if (texts.length == 0) { + return new float[0][]; + } + EmbedRequest[] pending = new EmbedRequest[texts.length]; + for (int i = 0; i < texts.length; i++) { + pending[i] = enqueue(task, texts[i]); + } + float[][] vectors = new float[texts.length][]; + for (int i = 0; i < texts.length; i++) { + vectors[i] = pending[i].awaitResult(); + } + return vectors; + } + + private EmbedRequest enqueue(TaskType task, String text) throws IndexException { + EmbedRequest request = new EmbedRequest(task, text); + try { + queue.put(request); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw IndexException.wrap("Embedding interrupted for model '" + name + "'", e); + } + return request; + } + + @Override + public void close() { + for (int i = 0; i < workers.length; i++) { + try { + queue.put(EmbedRequest.SHUTDOWN); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + break; + } + } + if (workers != null) { + for (InferenceWorker worker : workers) { + worker.awaitStop(); + } + } + LOG.debug("Closed embedding model '{}'", name); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/mapping/FieldSchema.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/mapping/FieldSchema.java new file mode 100644 index 000000000000..065429fa5bc9 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/mapping/FieldSchema.java @@ -0,0 +1,34 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.mapping; + +public interface FieldSchema { + String name(); + + boolean store(); + + boolean filter(); + + record TextFieldSchema( + String name, SearchParams search, boolean store, boolean filter) + implements FieldSchema { + public TextFieldSchema(String name, SearchParams search) { + this(name, search, true, false); + } + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/mapping/IndexMapping.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/mapping/IndexMapping.java new file mode 100644 index 000000000000..063dd794aac2 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/mapping/IndexMapping.java @@ -0,0 +1,83 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.mapping; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hive.search.config.IndexConfig; +import org.apache.hive.search.config.IndexStoreConfig; +import org.apache.hive.search.config.InferenceConfig; +import org.apache.hive.search.config.SearchConfig; +import org.apache.lucene.analysis.Analyzer; +import org.apache.lucene.analysis.core.KeywordAnalyzer; +import org.apache.lucene.analysis.miscellaneous.PerFieldAnalyzerWrapper; +import org.apache.lucene.analysis.standard.StandardAnalyzer; + +public record IndexMapping( + String indexName, Configuration configuration, Map fields) { + + public IndexConfig config() { + return new IndexConfig(configuration); + } + + public IndexStoreConfig store() { + return new IndexStoreConfig(configuration, indexName); + } + + public InferenceConfig inference() { + return new InferenceConfig(configuration); + } + + public SearchConfig search() { + return new SearchConfig(configuration); + } + + public FieldSchema fieldSchema(String fieldName) { + return fields.get(fieldName); + } + + public List hybridFields() { + List result = new ArrayList<>(); + for (Map.Entry entry : fields.entrySet()) { + if (entry.getValue() instanceof FieldSchema.TextFieldSchema text && text.search().hybrid()) { + result.add(entry.getKey()); + } + } + return result; + } + + public Optional soleHybridField() { + List hybridFields = hybridFields(); + return hybridFields.size() == 1 ? Optional.of(hybridFields.getFirst()) : Optional.empty(); + } + + public Analyzer analyzer() { + Map fieldAnalyzers = new HashMap<>(); + for (Map.Entry entry : fields.entrySet()) { + if (entry.getValue() instanceof FieldSchema.TextFieldSchema text && text.search().lexical()) { + fieldAnalyzers.put(entry.getKey(), new StandardAnalyzer()); + } + } + return new PerFieldAnalyzerWrapper(new KeywordAnalyzer(), fieldAnalyzers); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/mapping/SearchParams.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/mapping/SearchParams.java new file mode 100644 index 000000000000..89c324320c34 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/mapping/SearchParams.java @@ -0,0 +1,40 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.mapping; + +import org.apache.commons.lang3.StringUtils; + +public record SearchParams(boolean lexical, String semanticModel, VectorDistance distance) { + + public static SearchParams disabled() { + return new SearchParams(false, null, VectorDistance.COSINE); + } + + public boolean semantic() { + return StringUtils.isNotEmpty(semanticModel); + } + + public boolean hybrid() { + return lexical && semantic(); + } + + public enum VectorDistance { + COSINE, + DOT + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/mapping/TableDocument.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/mapping/TableDocument.java new file mode 100644 index 000000000000..860f60efcc13 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/mapping/TableDocument.java @@ -0,0 +1,117 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.mapping; + +import java.util.ArrayList; +import java.util.List; + +import org.apache.hive.search.exception.IndexException; +import org.apache.hive.search.mapping.field.Field; +import org.apache.hive.search.mapping.field.IdField; +import org.apache.hive.search.mapping.field.TextField; +import org.apache.lucene.document.BinaryDocValuesField; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StoredField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.util.BytesRef; + +public class TableDocument { + private static final int MAX_FIELD_SEARCH_SIZE = 32_768; + + private final List fields; + private final Document document; + private final IndexMapping indexMapping; + private final IdField idField; + public static final String FILTER_SUFFIX = ".filter"; + + public TableDocument(IdField idf, List otherFields, IndexMapping mapping) { + idField = idf; + indexMapping = mapping; + fields = new ArrayList<>(1 + otherFields.size()); + fields.add(idField); + fields.addAll(otherFields); + document = new Document(); + } + + public void fill(TextField field, FieldSchema.TextFieldSchema schema) + throws IndexException { + String trimmed = trim(field.value(), MAX_FIELD_SEARCH_SIZE); + + if (schema.filter()) { + document.add( + new StringField(field.name() + FILTER_SUFFIX, field.value(), + org.apache.lucene.document.Field.Store.NO)); + } + if (schema.search().lexical()) { + org.apache.lucene.document.Field.Store store = + schema.store() ? org.apache.lucene.document.Field.Store.YES + : org.apache.lucene.document.Field.Store.NO; + document.add(new org.apache.lucene.document.TextField(field.name(), trimmed, store)); + } else if (schema.store()) { + document.add(new StoredField(field.name(), field.value())); + } + if (schema.search().semantic()) { + if (field.embedding() == null) { + throw new IndexException("semantic field '" + field.name() + "' requires embedding"); + } + VectorSimilarityFunction similarity = schema.search().distance() == SearchParams.VectorDistance.DOT ? + VectorSimilarityFunction.DOT_PRODUCT : VectorSimilarityFunction.COSINE; + document.add(new KnnFloatVectorField(field.name(), field.embedding(), similarity)); + } + } + + private static String trim(String value, int max) { + return value.length() > max ? value.substring(0, max) : value; + } + + public List toDocuments() throws IndexException { + String id = idField().value(); + document.add(new BinaryDocValuesField("_id" + FILTER_SUFFIX, new BytesRef(id))); + document.add(new StoredField("_id", id)); + document.add(new StringField("_id" + FILTER_SUFFIX, id, + org.apache.lucene.document.Field.Store.NO)); + for (Field field : fields) { + if (field instanceof IdField) { + continue; + } + FieldSchema schema = indexMapping.fieldSchema(field.name()); + if (schema == null) { + continue; + } + if (schema instanceof FieldSchema.TextFieldSchema text + && field instanceof TextField tf) { + fill(tf, text); + } + } + return List.of(document); + } + + public void appendField(Field field) { + fields.add(field); + } + + public IdField idField() { + return idField; + } + + public List fields() { + return fields; + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/mapping/field/Field.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/mapping/field/Field.java new file mode 100644 index 000000000000..35ac9bc2e0de --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/mapping/field/Field.java @@ -0,0 +1,23 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.mapping.field; + +/** A single field value within a {@link org.apache.hive.search.mapping.TableDocument}. */ +public interface Field { + String name(); +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/mapping/field/IdField.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/mapping/field/IdField.java new file mode 100644 index 000000000000..5fdcfcded2c6 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/mapping/field/IdField.java @@ -0,0 +1,20 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.mapping.field; + +public record IdField(String name, String value) implements Field {} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/mapping/field/TextField.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/mapping/field/TextField.java new file mode 100644 index 000000000000..3eac83620d5c --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/mapping/field/TextField.java @@ -0,0 +1,49 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.mapping.field; + +import java.util.Objects; + +public record TextField(String name, String value, float[] embedding) implements Field { + public TextField(String name, String value) { + this(name, value, null); + } + + public TextField { + if ("_id".equals(name)) { + throw new IllegalArgumentException("use IdField for _id, not TextField"); + } + } + + public TextField withEmbedding(float[] newEmbedding) { + return new TextField(name, value, newEmbedding); + } + + @Override + public boolean equals(Object o) { + if (o == null || getClass() != o.getClass()) + return false; + TextField textField = (TextField) o; + return Objects.equals(name, textField.name) && Objects.equals(value, textField.value); + } + + @Override + public int hashCode() { + return Objects.hash(name, value); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/BootstrapIndexer.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/BootstrapIndexer.java new file mode 100644 index 000000000000..304dc0a64bdd --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/BootstrapIndexer.java @@ -0,0 +1,277 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.metastore; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.metastore.IMetaStoreClient; +import org.apache.hadoop.hive.metastore.RetryingMetaStoreClient; +import org.apache.hadoop.hive.metastore.api.Table; +import org.apache.hive.search.mapping.IndexMapping; +import org.apache.hive.search.mapping.TableDocument; +import org.apache.hive.search.exception.IndexException; +import org.apache.hive.search.config.IndexConfig; +import org.apache.hive.search.index.Indexer; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +final class BootstrapIndexer { + private static final Logger LOG = LoggerFactory.getLogger(BootstrapIndexer.class); + private static final List END_OF_STREAM = List.of(); + private static final TableBatch END_OF_WORK = new TableBatch(null, List.of()); + + private final Configuration configuration; + private final IndexConfig indexConfig; + private final IndexMapping mapping; + private final Indexer indexer; + private final IMetaStoreClient client; + private final boolean shareFetchClient; + + /** Package-private for unit tests: fetch workers reuse the injected client. */ + BootstrapIndexer(Configuration configuration, + IndexMapping mapping, Indexer indexer, + IMetaStoreClient client, boolean shareFetchClient) { + this.configuration = configuration; + this.indexConfig = new IndexConfig(configuration); + this.mapping = mapping; + this.indexer = indexer; + this.client = client; + this.shareFetchClient = shareFetchClient; + } + + void run(long notificationId) throws Exception { + long start = System.currentTimeMillis(); + List databases = client.getAllDatabases(); + BatchPlan plan = planBatches(databases); + LOG.info("Bootstrap planned {} table(s) in {} batch(es) across {} database(s)", + plan.expectedTableCount(), plan.batches().size(), databases.size()); + + int fetchThreads = indexConfig.getBootstrapFetchThreads(); + BlockingQueue workQueue = + new ArrayBlockingQueue<>(Math.max(fetchThreads * 2, 8)); + BlockingQueue> indexQueue = + new ArrayBlockingQueue<>(indexConfig.getBootstrapQueueDepth()); + AtomicReference failure = new AtomicReference<>(); + AtomicLong indexedTables = new AtomicLong(); + AtomicInteger completedBatches = new AtomicInteger(); + CountDownLatch fetchDone = new CountDownLatch(plan.batches().size()); + + Thread indexThread = startIndexConsumer(indexQueue, failure, indexedTables, completedBatches); + ExecutorService fetchPool = Executors.newFixedThreadPool(fetchThreads, r -> { + Thread thread = new Thread(r, "Index-Bootstrap-Fetch"); + thread.setDaemon(true); + return thread; + }); + try { + for (int i = 0; i < fetchThreads; i++) { + fetchPool.submit(() -> fetchTableWorker(plan, workQueue, indexQueue, failure, fetchDone)); + } + enqueueBatches(plan.batches(), workQueue, failure, fetchDone); + awaitFetchCompletion(fetchDone, failure); + indexQueue.put(END_OF_STREAM); + indexThread.join(); + rethrowFailure(failure); + validateBootstrapTableCount(plan.expectedTableCount().get()); + indexer.flush(notificationId, true); + indexer.syncBackup(); + long elapsed = System.currentTimeMillis() - start; + LOG.info("Built index for {} tables in {} batch(es) over {}ms", + plan.expectedTableCount(), plan.batches().size(), elapsed); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IndexException("Bootstrap indexing interrupted", e); + } finally { + workQueue.offer(END_OF_WORK); + fetchPool.shutdownNow(); + indexThread.interrupt(); + indexThread.join(TimeUnit.SECONDS.toMillis(30)); + } + } + + private void enqueueBatches(List batches, BlockingQueue workQueue, + AtomicReference failure, CountDownLatch fetchDone) throws InterruptedException { + int enqueued = 0; + try { + for (TableBatch batch : batches) { + if (failure.get() != null) { + break; + } + workQueue.put(batch); + enqueued++; + } + } finally { + for (int i = enqueued; i < batches.size(); i++) { + fetchDone.countDown(); + } + workQueue.put(END_OF_WORK); + } + } + + private BatchPlan planBatches(List databases) throws Exception { + int batchSize = indexConfig.getBootstrapBatchSize(); + List planned = new ArrayList<>(); + for (String db : databases) { + List tableNames = client.getAllTables(db); + for (int idx = 0; idx < tableNames.size(); idx += batchSize) { + int end = Math.min(idx + batchSize, tableNames.size()); + planned.add(new TableBatch(db, List.copyOf(tableNames.subList(idx, end)))); + } + } + return new BatchPlan(new AtomicInteger(0), planned); + } + + private void fetchTableWorker(BatchPlan plan, BlockingQueue workQueue, + BlockingQueue> indexQueue, AtomicReference failure, + CountDownLatch fetchDone) { + if (shareFetchClient) { + runFetchLoop(client, plan, workQueue, indexQueue, failure, fetchDone); + return; + } + try (IMetaStoreClient fetchClient = RetryingMetaStoreClient.getProxy(configuration, true)) { + runFetchLoop(fetchClient, plan, workQueue, indexQueue, failure, fetchDone); + } catch (Exception e) { + recordFailure(failure, e); + } + } + + private void runFetchLoop(IMetaStoreClient fetchClient, BatchPlan plan, + BlockingQueue workQueue, BlockingQueue> indexQueue, + AtomicReference failure, CountDownLatch fetchDone) { + try { + while (true) { + TableBatch batch = workQueue.take(); + if (batch == END_OF_WORK) { + workQueue.offer(END_OF_WORK); + return; + } + try { + if (failure.get() == null) { + List tables = + fetchClient.getTableObjectsByName(batch.database(), batch.tableNames()); + plan.expectedTableCount().addAndGet(tables.size()); + List documents = new ArrayList<>(tables.size()); + for (Table table : tables) { + documents.add(MetastoreTableMapper.fromTable(table, mapping)); + } + indexQueue.put(documents); + } + } catch (Exception e) { + recordFailure(failure, e); + } finally { + fetchDone.countDown(); + } + } + } catch (Exception e) { + recordFailure(failure, e); + } + } + + private Thread startIndexConsumer(BlockingQueue> indexQueue, + AtomicReference failure, AtomicLong indexedTables, + AtomicInteger completedBatches) { + Thread indexThread = new Thread(() -> { + long lastProgressLog = System.currentTimeMillis(); + try { + while (true) { + List documents = indexQueue.take(); + if (documents == END_OF_STREAM) { + break; + } + if (failure.get() != null) { + break; + } + indexer.addDocuments(documents); + long indexed = indexedTables.addAndGet(documents.size()); + int done = completedBatches.incrementAndGet(); + indexer.flush(-1000000L, false); + long now = System.currentTimeMillis(); + if (now - lastProgressLog >= indexConfig.getBootstrapProgressIntervalMs()) { + LOG.info("Bootstrap progress: indexed {} table(s) in {} batch(es)", indexed, done); + lastProgressLog = now; + } + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + recordFailure(failure, e); + } catch (Exception e) { + recordFailure(failure, e); + } + }, "Index-Bootstrap-Writer"); + indexThread.setDaemon(true); + indexThread.start(); + return indexThread; + } + + private static void awaitFetchCompletion(CountDownLatch fetchDone, + AtomicReference failure) throws Exception { + while (fetchDone.getCount() > 0) { + if (failure.get() != null) { + break; + } + fetchDone.await(1, TimeUnit.SECONDS); + } + rethrowFailure(failure); + } + + private void validateBootstrapTableCount(int expectedTableCount) throws IndexException { + int indexed = indexer.writer().getDocStats().numDocs; + if (indexed != expectedTableCount) { + throw new IndexException( + "Bootstrap index doc count mismatch: indexed " + indexed + + " documents but metastore has " + expectedTableCount + " tables"); + } + } + + private static void recordFailure(AtomicReference failure, Exception error) { + failure.compareAndSet(null, error); + LOG.error("Bootstrap indexing failed", error); + } + + private static void rethrowFailure(AtomicReference failure) throws Exception { + Exception error = failure.get(); + if (error != null) { + if (error instanceof IndexException indexException) { + throw indexException; + } + if (error instanceof IOException ioException) { + throw ioException; + } + if (error instanceof InterruptedException interruptedException) { + Thread.currentThread().interrupt(); + throw interruptedException; + } + throw new IndexException("Bootstrap indexing failed", error); + } + } + + private record BatchPlan(AtomicInteger expectedTableCount, List batches) {} + + private record TableBatch(String database, List tableNames) {} +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/MetastoreCluster.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/MetastoreCluster.java new file mode 100644 index 000000000000..686e4c1e784e --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/MetastoreCluster.java @@ -0,0 +1,79 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.metastore; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +import org.apache.commons.lang3.tuple.Pair; +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.common.TableName; +import org.apache.hadoop.hive.metastore.leader.LeaderElection; +import org.apache.hadoop.hive.metastore.leader.LeaderException; +import org.apache.hadoop.hive.metastore.leader.LeaseLeaderElection; + +@SuppressWarnings("rawtypes, unchecked") +public class MetastoreCluster implements AutoCloseable { + private static final String CLUSTER_ELECTION_METHOD = "metastore.index.election.method"; + private static final Map> CACHED = new ConcurrentHashMap<>(); + + static { + try { + CACHED.put("default", Pair.of(new LeaseLeaderElection(), + new TableName("__INDEX_CATALOG__", "__INDEX_DATABASE__", "__INDEX_TABLE__"))); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + private final LeaderElection election; + + public MetastoreCluster(Configuration configuration, + LeaderElection.LeadershipStateListener... listeners) throws IOException, LeaderException { + String methodName = configuration.get(CLUSTER_ELECTION_METHOD, "default"); + Pair leo = CACHED.get(methodName); + if (leo == null) { + throw new IOException("No elector configured for " + methodName); + } + this.election = leo.getKey(); + if (listeners != null) { + Arrays.stream(listeners).forEach(this.election::addStateListener); + } + + this.election.tryBeLeader(configuration, leo.getValue()); + } + + boolean isLeader() { + return this.election.isLeader(); + } + + @Override + public void close() throws Exception { + if (election != null) { + election.close(); + } + } + + public static void injectElection(Configuration configuration, + String methodName, LeaderElection election, T mutex) { + configuration.set(CLUSTER_ELECTION_METHOD, methodName); + CACHED.put(methodName, Pair.of(election, mutex)); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/MetastoreEventHandler.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/MetastoreEventHandler.java new file mode 100644 index 000000000000..d92f93ebad0d --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/MetastoreEventHandler.java @@ -0,0 +1,385 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.metastore; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import java.util.concurrent.atomic.AtomicBoolean; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.common.DatabaseName; +import org.apache.hadoop.hive.common.TableName; +import org.apache.hadoop.hive.metastore.IMetaStoreClient; +import org.apache.hadoop.hive.metastore.RetryingMetaStoreClient; +import org.apache.hadoop.hive.metastore.api.NotificationEvent; +import org.apache.hadoop.hive.metastore.api.NotificationEventRequest; +import org.apache.hadoop.hive.metastore.api.NotificationEventResponse; +import org.apache.hadoop.hive.metastore.api.Table; +import org.apache.hadoop.hive.metastore.messaging.AlterTableMessage; +import org.apache.hadoop.hive.metastore.messaging.CreateTableMessage; +import org.apache.hadoop.hive.metastore.messaging.MessageBuilder; +import org.apache.hadoop.hive.metastore.messaging.MessageDeserializer; +import org.apache.hadoop.hive.metastore.messaging.MessageFactory; +import org.apache.hive.search.exception.IndexNotHealthyException; +import org.apache.hive.search.config.IndexConfig; +import org.apache.thrift.TException; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Polls HMS notification events, coalesces them into batches, and dispatches index mutations. + * + *

On repeated batch failure the handler falls back to single-event apply so one poison event + * does not block the entire batch forever. If a poison event still cannot be applied, the index + * is marked unhealthy and search is blocked until the root cause is fixed. + */ +public class MetastoreEventHandler implements AutoCloseable { + private static final Logger LOG = LoggerFactory.getLogger(MetastoreEventHandler.class); + private final AtomicBoolean stopped = new AtomicBoolean(false); + private final List listeners = Collections.synchronizedList(new ArrayList<>()); + private final IndexConfig indexConfig; + private final IMetaStoreClient client; + private final MessageDeserializer deserializer; + private Thread metaRefresher; + private long lastEventId; + + /** Consecutive batch failures at {@link #failedBatchStartId}. */ + private int consecutiveBatchFailures; + private long failedBatchStartId = -1; + + private MetastoreEventHandler(Configuration configuration) throws TException { + this(configuration, RetryingMetaStoreClient.getProxy(configuration, true)); + } + + /** Package-private for unit tests with a stub {@link IMetaStoreClient}. */ + MetastoreEventHandler(Configuration configuration, IMetaStoreClient client) { + Configuration conf = new Configuration(Objects.requireNonNull(configuration)); + this.indexConfig = new IndexConfig(conf); + this.client = Objects.requireNonNull(client); + this.deserializer = MessageFactory.getDefaultInstance(conf).getDeserializer(); + } + + public static MetastoreEventHandler of(Configuration conf, MetastoreEventListener... listeners) + throws TException { + MetastoreEventHandler catalogService = new MetastoreEventHandler(conf); + return catalogService.addListeners(listeners); + } + + public MetastoreEventHandler addListeners(MetastoreEventListener... listeners) { + return addListeners(Arrays.asList(listeners)); + } + + public MetastoreEventHandler addListeners(Collection listeners) { + this.listeners.addAll(listeners); + return this; + } + + public MetastoreEventHandler start(long nid) throws Exception { + this.lastEventId = nid; + final long sleepInterval = indexConfig.getPollNotificationInterval(); + final int pollMax = indexConfig.getEventPollMax(); + final long unhealthyBackoffMs = indexConfig.getEventUnhealthyBackoffMs(); + this.metaRefresher = new Thread(() -> { + while (!stopped.get()) { + try { + int processed = getNextMetastoreEvents(pollMax); + if (processed <= 0) { + Thread.sleep(sleepInterval); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + break; + } catch (IndexNotHealthyException e) { + LOG.error("Incremental index update halted: {}", e.getMessage(), e); + notifyListenersStatus(false, e); + try { + Thread.sleep(unhealthyBackoffMs); + } catch (InterruptedException ex) { + Thread.currentThread().interrupt(); + break; + } + } + } + }); + this.metaRefresher.setName("Metastore-Event-Poller"); + this.metaRefresher.start(); + return this; + } + + @Override + public void close() throws Exception { + stopped.set(true); + if (metaRefresher != null) { + metaRefresher.interrupt(); + metaRefresher.join(); + } + if (client != null) { + client.close(); + } + } + + public int getNextMetastoreEvents(int eventCount) throws IndexNotHealthyException { + try { + NotificationEventRequest request = new NotificationEventRequest(lastEventId); + request.setMaxEvents(eventCount); + NotificationEventResponse resp = client.getNextNotification(request, true, null); + if (resp == null || resp.getEvents() == null || resp.getEvents().isEmpty()) { + LOG.debug("No event found since the last event id: {}", lastEventId); + return 0; + } + List events = resp.getEvents(); + long batchStartId = events.get(0).getEventId(); + + MetastoreEventListener.IndexTask task; + try { + task = buildTask(events); + } catch (Exception parseError) { + LOG.warn( + "Failed to build notification batch starting at event {}; falling back to single-event apply", + batchStartId, + parseError); + int applied = applyEventsIndividually(events); + if (applied > 0) { + notifyListenersStatus(true); + resetBatchFailureState(); + return applied; + } + backoffAfterFailure(); + return 0; + } + try { + notifyListeners(task); + lastEventId = task.lastEventId; + notifyListenersStatus(true); + resetBatchFailureState(); + return events.size(); + } catch (Exception batchError) { + recordBatchFailure(batchStartId, task, batchError); + if (consecutiveBatchFailures >= indexConfig.getEventBatchMaxFailures()) { + LOG.warn( + "Batch apply failed {} time(s) at event {}; falling back to single-event apply", + consecutiveBatchFailures, + failedBatchStartId); + int applied = applyEventsIndividually(events); + if (applied > 0) { + notifyListenersStatus(true); + resetBatchFailureState(); + return applied; + } + } + backoffAfterFailure(); + return 0; + } + } catch (IndexNotHealthyException e) { + throw e; + } catch (Exception e) { + LOG.warn("Failed to fetch or parse notification events after lastEventId={}", lastEventId, e); + backoffAfterFailure(); + } + return 0; + } + + private int applyEventsIndividually(List events) + throws IndexNotHealthyException { + int applied = 0; + for (NotificationEvent event : events) { + if (lastEventId >= event.getEventId()) { + continue; + } + MetastoreEventListener.IndexTask task = + new MetastoreEventListener.IndexTask(); + task.firstEventId = event.getEventId(); + try { + dispatchEvent(event, task); + task.lastEventId = event.getEventId(); + notifyListeners(task); + lastEventId = event.getEventId(); + applied++; + } catch (Exception e) { + long poisonEventId = event.getEventId(); + String poisonEventType = event.getEventType(); + if (indexConfig.isEventSkipPoison()) { + LOG.error( + "Skipping poison notification event {} (type={}); index may be stale for this object", + poisonEventId, + poisonEventType, + e); + lastEventId = event.getEventId(); + applied++; + } else { + LOG.error( + "Stuck on notification event {} (type={}); committed lastEventId={}", + poisonEventId, + poisonEventType, + lastEventId, + e); + String progress = applied > 0 + ? applied + " prior event(s) in batch were applied; " + : ""; + throw new IndexNotHealthyException( + "Cannot apply notification event " + + poisonEventId + + " (type=" + + poisonEventType + + "); " + + progress + + "committed lastEventId=" + + lastEventId + + ". Fix root cause, set " + + IndexConfig.EVENT_SKIP_POISON + + "=true to skip, or rebuild the index.", + e); + } + } + } + if (applied > 0) { + LOG.info("Applied {} notification event(s) individually; lastEventId={}", applied, lastEventId); + } + return applied; + } + + private MetastoreEventListener.IndexTask buildTask(List events) + throws Exception { + MetastoreEventListener.IndexTask task = new MetastoreEventListener.IndexTask(); + for (int i = 0; i < events.size(); i++) { + NotificationEvent event = events.get(i); + if (lastEventId >= event.getEventId()) { + throw new IllegalStateException( + "Out-of-order metastore notification event: lastEventId=" + lastEventId + + ", eventId=" + event.getEventId()); + } + if (i == 0) { + task.firstEventId = event.getEventId(); + } + dispatchEvent(event, task); + task.lastEventId = event.getEventId(); + } + return task; + } + + private void notifyListeners(MetastoreEventListener.IndexTask task) throws Exception { + for (MetastoreEventListener listener : listeners) { + listener.notifyIndexTask(task); + } + } + + private void notifyListenersStatus(boolean healthy, IndexNotHealthyException... errors) { + for (MetastoreEventListener listener : listeners) { + listener.notifyIndexState(healthy, errors); + } + } + + private void recordBatchFailure( + long batchStartId, MetastoreEventListener.IndexTask task, Exception e) { + if (failedBatchStartId != batchStartId) { + failedBatchStartId = batchStartId; + consecutiveBatchFailures = 1; + } else { + consecutiveBatchFailures++; + } + LOG.warn( + "Failed to apply notification batch (attempt {}/{}, committed lastEventId={}, " + + "batch event range {}-{}, upserts={}, tableDrops={}, dbDrops={})", + consecutiveBatchFailures, + indexConfig.getEventBatchMaxFailures(), + lastEventId, + task.firstEventId, + task.lastEventId, + task.tablesToAdd.size(), + task.tablesToDrop.size(), + task.databasesToDrop.size(), + e); + } + + private void resetBatchFailureState() { + consecutiveBatchFailures = 0; + failedBatchStartId = -1; + } + + private void backoffAfterFailure() { + long backoffMs = indexConfig.getEventFailureBackoffMs(); + if (backoffMs <= 0) { + return; + } + try { + Thread.sleep(backoffMs); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + + private void dispatchEvent( + NotificationEvent event, + MetastoreEventListener.IndexTask task) + throws Exception { + String message = event.getMessage(); + switch (event.getEventType()) { + case MessageBuilder.CREATE_TABLE_EVENT -> { + CreateTableMessage createTableMessage = deserializer.getCreateTableMessage(message); + Table table = createTableMessage.getTableObj(); + TableName tableName = new TableName(table.getCatName(), table.getDbName(), table.getTableName()); + task.tablesToAdd.put(tableName, table); + } + case MessageBuilder.DROP_TABLE_EVENT -> { + deserializer.getDropTableMessage(message); + TableName tableName = new TableName(event.getCatName(), event.getDbName(), event.getTableName()); + if (task.tablesToAdd.containsKey(tableName)) { + task.tablesToAdd.remove(tableName); + } else { + task.tablesToDrop.add(tableName); + } + } + case MessageBuilder.DROP_DATABASE_EVENT -> { + DatabaseName databaseName = new DatabaseName(event.getCatName(), event.getDbName()); + task.databasesToDrop.add(databaseName); + for (TableName tableName : new ArrayList<>(task.tablesToAdd.keySet())) { + if (databaseName.equals(new DatabaseName(tableName.getCat(), tableName.getDb()))) { + task.tablesToAdd.remove(tableName); + } + } + task.tablesToDrop.removeIf( + tableName -> databaseName.equals(new DatabaseName(tableName.getCat(), tableName.getDb()))); + } + case MessageBuilder.ALTER_TABLE_EVENT -> { + AlterTableMessage alterTableMessage = deserializer.getAlterTableMessage(message); + Table tableAfter = alterTableMessage.getTableObjAfter(); + Table tableBefore = alterTableMessage.getTableObjBefore(); + if (!MetastoreTableMapper.hasIndexedFieldsChanged(tableBefore, tableAfter)) { + break; + } + TableName tblNameBefore = new TableName(tableBefore.getCatName(), + tableBefore.getDbName(), tableBefore.getTableName()); + TableName tblNameAfter = new TableName(tableAfter.getCatName(), + tableAfter.getDbName(), tableAfter.getTableName()); + if (task.tablesToAdd.containsKey(tblNameBefore)) { + task.tablesToAdd.remove(tblNameBefore); + } else { + task.tablesToDrop.add(tblNameBefore); + } + task.tablesToAdd.put(tblNameAfter, tableAfter); + } + default -> { + // Ignored event types still advance the notification cursor when applied individually. + } + } + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/MetastoreEventListener.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/MetastoreEventListener.java new file mode 100644 index 000000000000..dcc8a04787a2 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/MetastoreEventListener.java @@ -0,0 +1,52 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.metastore; + +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.Map; +import java.util.Set; + +import org.apache.hadoop.hive.common.DatabaseName; +import org.apache.hadoop.hive.common.TableName; +import org.apache.hadoop.hive.metastore.api.Table; +import org.apache.hive.search.exception.IndexException; +import org.apache.hive.search.exception.IndexNotHealthyException; + +public interface MetastoreEventListener { + + /** + * Apply a coalesced batch of index mutations. Implementations must either complete all + * mutations or throw; partial success must remain safe to retry idempotently. + */ + default void notifyIndexTask(IndexTask task) throws IndexException, java.io.IOException { + + } + + default void notifyIndexState(boolean healthy, IndexNotHealthyException... e) { + + } + + class IndexTask { + public long firstEventId; + public long lastEventId; + public Set tablesToDrop = new LinkedHashSet<>(); + public Set databasesToDrop = new LinkedHashSet<>(); + public Map tablesToAdd = new LinkedHashMap<>(); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/MetastoreIndexer.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/MetastoreIndexer.java new file mode 100644 index 000000000000..ea5223d4053a --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/MetastoreIndexer.java @@ -0,0 +1,356 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.metastore; + +import java.io.IOException; +import java.util.List; +import java.util.Optional; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.common.DatabaseName; +import org.apache.hadoop.hive.metastore.IMetaStoreClient; +import org.apache.hadoop.hive.metastore.RetryingMetaStoreClient; +import org.apache.hadoop.hive.metastore.api.NotificationEventRequest; +import org.apache.hadoop.hive.metastore.leader.LeaderElection; +import org.apache.hive.search.mapping.TableDocument; +import org.apache.hive.search.exception.IndexException; +import org.apache.hive.search.exception.IndexNotHealthyException; +import org.apache.hive.search.config.IndexConfig; +import org.apache.hive.search.index.Indexer; +import org.apache.hive.search.index.IndexManager; +import org.apache.hive.search.index.manifest.IndexManifest; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public final class MetastoreIndexer implements AutoCloseable { + private static final Logger LOG = LoggerFactory.getLogger(MetastoreIndexer.class); + + private final IMetaStoreClient client; + private final long lastEventId; + private final Indexer indexer; + private final MetastoreCluster cluster; + private final MetastoreEventHandler handler; + private final IndexManager indexManager; + private final FlushIndexListener flushIndexListener; + private final boolean shareBootstrapFetchClient; + + public MetastoreIndexer(Configuration configuration, IndexManager indexManager, Indexer indexer) + throws Exception { + this(configuration, indexManager, indexer, + RetryingMetaStoreClient.getProxy(configuration, true), false); + } + + MetastoreIndexer(Configuration configuration, IndexManager indexManager, Indexer indexer, + IMetaStoreClient client) throws Exception { + this(configuration, indexManager, indexer, client, true); + } + + MetastoreIndexer(Configuration configuration, IndexManager indexManager, Indexer indexer, + IMetaStoreClient client, boolean shareBootstrapFetchClient) throws Exception { + this.client = client; + this.shareBootstrapFetchClient = shareBootstrapFetchClient; + this.indexer = indexer; + this.indexManager = indexManager; + this.flushIndexListener = new FlushIndexListener(configuration); + this.handler = new MetastoreEventHandler(configuration, client); + this.handler.addListeners(flushIndexListener, indexManager); + this.cluster = new MetastoreCluster(configuration, flushIndexListener); + this.lastEventId = initialize(); + this.indexManager.setIndexedNid(lastEventId); + } + + private boolean rebuildIndex() throws Exception { + indexManager.resolveInterruptedRestore(); + if (isIndexValid(indexManager.readLocalManifest().orElse(null))) { + return false; + } + if (isIndexValid(indexManager.readRemoteManifest().orElse(null))) { + if (indexManager.restoreBackup()) { + return false; + } + LOG.warn("Failed to restore index from remote backup; clearing partial local index"); + indexManager.clearLocalIndex(); + return true; + } + indexManager.clearLocalIndex(); + return true; + } + + private boolean isIndexValid(IndexManifest indexManifest) + throws Exception { + if (indexManifest == null) { + return false; + } + String currentModelName = indexManager.mapping().inference().modelName(); + String modelName = indexManifest.modelName(); + if (!currentModelName.equals(modelName)) { + return false; + } + long notificationId = indexManifest.lastEventId(); + return canCatchUp(notificationId); + } + + private boolean canCatchUp(long notificationId) throws Exception { + if (notificationId < 0) { + return false; + } + try { + NotificationEventRequest request = new NotificationEventRequest(notificationId); + request.setMaxEvents(10); + client.getNextNotification(request, false, null); + return true; + } catch (IllegalStateException ignored) { + LOG.debug("Current index lags too behind, will rebuild"); + } + return false; + } + + private long initialize() throws Exception { + boolean rebuild = rebuildIndex(); + if (rebuild && !cluster.isLeader()) { + long notificationId = waitForRemoteBackup(); + indexer.initialize(); + if (notificationId >= 0) { + return notificationId; + } + return rebuildLeaderIndex(); + } + indexer.initialize(); + if (rebuild) { + return rebuildLeaderIndex(); + } + return indexManager.readLocalManifest().get().lastEventId(); + } + + private long rebuildLeaderIndex() throws Exception { + indexManager.clearRemoteIndex(); + long notificationId = client.getCurrentNotificationEventId().getEventId(); + new BootstrapIndexer( + indexManager.mapping().configuration(), + indexManager.mapping(), + indexer, + client, + shareBootstrapFetchClient).run(notificationId); + return notificationId; + } + + public void start() throws Exception { + handler.start(lastEventId); + flushIndexListener.start(lastEventId); + } + + /** Package-private for integration tests. */ + int pollEvents(int count) throws IndexNotHealthyException { + return handler.getNextMetastoreEvents(count); + } + + /** Package-private for integration tests. */ + void flushCheckpoint() throws IOException { + try { + indexer.flush(client.getCurrentNotificationEventId().getEventId(), true); + } catch (Exception e) { + if (e instanceof IOException ioException) { + throw ioException; + } + throw new IOException("Failed to flush index checkpoint", e); + } + } + + /** Package-private for integration tests. */ + void syncBackup() throws IOException { + indexer.syncBackup(); + } + + /** + * Wait for a follower-ready remote backup and restore it before the IndexWriter is opened. + * + * @return restored notification id, or -1 if this instance became leader while waiting + */ + private long waitForRemoteBackup() throws Exception { + String currentModelName = indexManager.mapping().inference().modelName(); + for (; !cluster.isLeader(); Thread.sleep(3000)) { + Optional indexManifest = indexManager.readRemoteManifest(); + if (indexManifest.isEmpty()) { + continue; + } + IndexManifest manifest = indexManifest.get(); + if (!currentModelName.equals(manifest.modelName())) { + LOG.debug("Remote index model {} does not match configured model {}, waiting", + manifest.modelName(), currentModelName); + continue; + } + long notificationId = manifest.lastEventId(); + if (notificationId > 0 && canCatchUp(notificationId)) { + if (indexManager.restoreBackup()) { + return notificationId; + } + LOG.warn("Failed to restore the index from remote directory, will retry"); + } + } + indexManager.clearLocalIndex(); + return -1; + } + + @Override + public void close() throws Exception { + try { + cluster.close(); + handler.close(); + } finally { + flushIndexListener.close(); + client.close(); + } + } + + private class FlushIndexListener + implements MetastoreEventListener, + LeaderElection.LeadershipStateListener, + AutoCloseable { + private volatile long eventId; + private volatile long lastCommittedEventId; + private volatile boolean isLeader; + private volatile Thread replicateThread; + private volatile boolean started = false; + private final Thread commitThread; + private final IndexConfig indexConfig; + + public FlushIndexListener(Configuration configuration) { + this.indexConfig = new IndexConfig(configuration); + this.commitThread = getIndexCommitThread(); + } + + public void start(long initialEventId) { + this.eventId = initialEventId; + this.lastCommittedEventId = initialEventId; + this.commitThread.start(); + this.started = true; + } + + private Thread getIndexCommitThread() { + Thread commitThread = new Thread(() -> { + while (!Thread.currentThread().isInterrupted()) { + try { + Thread.sleep(indexConfig.getFlushInterval()); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + break; + } + if (eventId <= lastCommittedEventId) { + continue; + } + try { + if (indexer.flush(eventId, false)) { + lastCommittedEventId = eventId; + } else if (eventId - lastCommittedEventId > indexConfig.getForceFlushEventGap()) { + // Many events processed with no Lucene changes; advance checkpoint metadata. + if (indexer.flush(eventId, true)) { + lastCommittedEventId = eventId; + } + } + } catch (IOException e) { + LOG.warn("Error flushing the index", e); + } + } + }); + commitThread.setName("Index-Commit"); + commitThread.setDaemon(true); + return commitThread; + } + + private Thread getIndexReplicateThread() { + Thread replicateThread = new Thread(() -> { + while (isLeader && !Thread.currentThread().isInterrupted()) { + long interval = 10000; + if (started) { + try { + indexer.syncBackup(); + interval = indexConfig.getSyncInterval(); + } catch (IOException e) { + LOG.warn("Error replicating the index to remote directory", e); + } + } + try { + Thread.sleep(interval); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + break; + } + } + }); + replicateThread.setDaemon(true); + replicateThread.setName("Index-Replica"); + return replicateThread; + } + + @Override + public void takeLeadership(LeaderElection leaderElection) { + isLeader = true; + if (indexManager.hasBackup()) { + if (replicateThread != null && replicateThread.isAlive()) { + return; + } + replicateThread = getIndexReplicateThread(); + replicateThread.start(); + } + } + + @Override + public void lossLeadership(LeaderElection leaderElection) { + isLeader = false; + if (replicateThread != null) { + replicateThread.interrupt(); + replicateThread = null; + } + } + + @Override + public void notifyIndexTask(IndexTask task) throws IndexException, IOException { + if (!task.databasesToDrop.isEmpty()) { + indexer.deleteDatabases(task.databasesToDrop.toArray(new DatabaseName[0])); + } + if (!task.tablesToDrop.isEmpty()) { + String[] docIds = task.tablesToDrop.stream() + .map(MetastoreTableMapper::tableId).toList().toArray(new String[0]); + indexer.delete(docIds); + } + if (!task.tablesToAdd.isEmpty()) { + List newDocs = task.tablesToAdd.values().stream() + .map(t -> MetastoreTableMapper.fromTable(t, indexManager.mapping())).toList(); + indexer.addDocuments(newDocs); + } + eventId = task.lastEventId; + } + + @Override + public void close() throws Exception { + isLeader = false; + try { + if (replicateThread != null) { + replicateThread.interrupt(); + replicateThread = null; + } + commitThread.interrupt(); + } finally { + if (eventId > lastCommittedEventId) { + indexer.flush(eventId, true); + lastCommittedEventId = eventId; + } + } + } + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/MetastoreSchemas.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/MetastoreSchemas.java new file mode 100644 index 000000000000..040795309b6e --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/MetastoreSchemas.java @@ -0,0 +1,79 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.metastore; + +import java.util.LinkedHashMap; +import java.util.Map; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hive.search.mapping.FieldSchema; +import org.apache.hive.search.mapping.IndexMapping; +import org.apache.hive.search.mapping.SearchParams; + +public final class MetastoreSchemas { + + private MetastoreSchemas() {} + + public static IndexMapping defaultHiveTablesMapping(String indexName, + String semanticModel, Configuration conf) { + Map fields = new LinkedHashMap<>(); + fields.put(MetastoreTableMapper.FIELD_DB, filterText(MetastoreTableMapper.FIELD_DB)); + fields.put(MetastoreTableMapper.FIELD_TABLE, tableNameText(MetastoreTableMapper.FIELD_TABLE)); + fields.put(MetastoreTableMapper.FIELD_OWNER, filterText(MetastoreTableMapper.FIELD_OWNER)); + fields.put(MetastoreTableMapper.FIELD_TABLE_TYPE, filterText(MetastoreTableMapper.FIELD_TABLE_TYPE)); + fields.put(MetastoreTableMapper.FIELD_LOCATION, storedText(MetastoreTableMapper.FIELD_LOCATION)); + fields.put(MetastoreTableMapper.FIELD_COMMENT, lexicalText(MetastoreTableMapper.FIELD_COMMENT)); + fields.put(MetastoreTableMapper.FIELD_COLUMNS, storedText(MetastoreTableMapper.FIELD_COLUMNS)); + fields.put( + MetastoreTableMapper.FIELD_COLUMN_NAMES, + lexicalText(MetastoreTableMapper.FIELD_COLUMN_NAMES)); + fields.put( + MetastoreTableMapper.FIELD_COLUMN_COMMENTS, + lexicalText(MetastoreTableMapper.FIELD_COLUMN_COMMENTS)); + fields.put( + MetastoreTableMapper.FIELD_SEARCH_TEXT, + hybridText(MetastoreTableMapper.FIELD_SEARCH_TEXT, semanticModel)); + return new IndexMapping(indexName, conf, fields); + } + + private static FieldSchema.TextFieldSchema filterText(String name) { + return new FieldSchema.TextFieldSchema(name, SearchParams.disabled(), true, true); + } + + private static FieldSchema.TextFieldSchema storedText(String name) { + return new FieldSchema.TextFieldSchema(name, SearchParams.disabled(), true, false); + } + + private static FieldSchema.TextFieldSchema lexicalText(String name) { + return new FieldSchema.TextFieldSchema( + name, new SearchParams(true, null, SearchParams.VectorDistance.COSINE), true, false); + } + + private static FieldSchema.TextFieldSchema tableNameText(String name) { + return new FieldSchema.TextFieldSchema( + name, new SearchParams(true, null, SearchParams.VectorDistance.COSINE), true, true); + } + + private static FieldSchema.TextFieldSchema hybridText(String name, String semanticModel) { + return new FieldSchema.TextFieldSchema( + name, + new SearchParams(true, semanticModel, SearchParams.VectorDistance.COSINE), + false, + false); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/MetastoreTableMapper.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/MetastoreTableMapper.java new file mode 100644 index 000000000000..512a787f5e6e --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/MetastoreTableMapper.java @@ -0,0 +1,194 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.metastore; + +import java.util.ArrayList; +import java.util.Locale; +import java.util.List; +import java.util.Objects; + +import org.apache.commons.lang3.StringUtils; +import org.apache.hadoop.hive.common.TableName; +import org.apache.hadoop.hive.metastore.api.FieldSchema; +import org.apache.hadoop.hive.metastore.api.Table; +import org.apache.hive.search.mapping.IndexMapping; +import org.apache.hive.search.mapping.TableDocument; +import org.apache.hive.search.mapping.field.Field; +import org.apache.hive.search.mapping.field.IdField; +import org.apache.hive.search.mapping.field.TextField; + +public final class MetastoreTableMapper { + public static final String FIELD_DB = "db"; + public static final String FIELD_TABLE = "table_name"; + public static final String FIELD_OWNER = "owner"; + public static final String FIELD_TABLE_TYPE = "table_type"; + public static final String FIELD_LOCATION = "location"; + public static final String FIELD_COMMENT = "comment"; + public static final String FIELD_COLUMNS = "columns"; + public static final String FIELD_COLUMN_NAMES = "column_names"; + public static final String FIELD_COLUMN_COMMENTS = "column_comments"; + public static final String FIELD_SEARCH_TEXT = "search_text"; + /** Relative boosts for {@code table_keyword} ranking: table name > column name > comment. */ + public static final float KEYWORD_BOOST_TABLE_NAME = 3.0f; + public static final float KEYWORD_BOOST_COLUMN_NAME = 2.0f; + public static final float KEYWORD_BOOST_COMMENT = 1.0f; + /** Lexical fields used for table keyword search, highest boost first. */ + public static final List KEYWORD_SEARCH_FIELDS = List.of( + new KeywordSearchField(FIELD_TABLE, KEYWORD_BOOST_TABLE_NAME), + new KeywordSearchField(FIELD_COLUMN_NAMES, KEYWORD_BOOST_COLUMN_NAME), + new KeywordSearchField(FIELD_COMMENT, KEYWORD_BOOST_COMMENT), + new KeywordSearchField(FIELD_COLUMN_COMMENTS, KEYWORD_BOOST_COMMENT)); + + public record KeywordSearchField(String field, float boost) {} + + private MetastoreTableMapper() {} + + public static String tableId(String catalog, String db, String table) { + return catalog + "." + db + "." + table; + } + + public static String tableId(TableName tableName) { + return tableId(tableName.getCat(), tableName.getDb(), tableName.getTable()); + } + + public static TableDocument fromTable(Table table, IndexMapping indexMapping) { + String db = table.getDbName(); + String name = table.getTableName(); + String catalog = table.getCatName(); + String id = tableId(catalog, db, name); + String owner = nullToEmpty(table.getOwner()); + String tableType = nullToEmpty(table.getTableType()); + String location = tableLocation(table); + String comment = tableComment(table); + String columns = formatColumnsForStorage(table); + String columnNames = formatColumnNamesForSearch(table); + String columnComments = formatColumnCommentsForSearch(table); + String searchText = buildSearchText(name, comment, table); + + List fields = new ArrayList<>(10); + fields.add(new TextField(FIELD_DB, db)); + fields.add(new TextField(FIELD_TABLE, name.toLowerCase(Locale.ROOT))); + fields.add(new TextField(FIELD_OWNER, owner)); + fields.add(new TextField(FIELD_TABLE_TYPE, tableType)); + fields.add(new TextField(FIELD_LOCATION, location)); + fields.add(new TextField(FIELD_COMMENT, comment)); + fields.add(new TextField(FIELD_COLUMNS, columns)); + fields.add(new TextField(FIELD_COLUMN_NAMES, columnNames)); + fields.add(new TextField(FIELD_COLUMN_COMMENTS, columnComments)); + fields.add(new TextField(FIELD_SEARCH_TEXT, searchText)); + return new TableDocument(new IdField("_id", id), fields, indexMapping); + } + + /** Returns true when an alter would change any value written to the search index. */ + public static boolean hasIndexedFieldsChanged(Table before, Table after) { + Objects.requireNonNull(before); + Objects.requireNonNull(after); + if (!Objects.equals(before.getCatName(), after.getCatName()) + || !Objects.equals(before.getDbName(), after.getDbName()) + || !Objects.equals(before.getTableName(), after.getTableName())) { + return true; + } + return !nullToEmpty(before.getOwner()).equals(nullToEmpty(after.getOwner())) + || !nullToEmpty(before.getTableType()).equals(nullToEmpty(after.getTableType())) + || !tableLocation(before).equals(tableLocation(after)) + || !tableComment(before).equals(tableComment(after)) + || !formatColumnsForStorage(before).equals(formatColumnsForStorage(after)) + || !formatColumnNamesForSearch(before).equals(formatColumnNamesForSearch(after)) + || !formatColumnCommentsForSearch(before).equals(formatColumnCommentsForSearch(after)); + } + + private static String tableLocation(Table table) { + return table.getSd() == null ? "" : nullToEmpty(table.getSd().getLocation()); + } + + private static String tableComment(Table table) { + if (table.getParameters() == null || table.getParameters().get("comment") == null) { + return ""; + } + return nullToEmpty(table.getParameters().get("comment")); + } + + private static String buildSearchText(String tableName, String comment, Table table) { + List parts = new ArrayList<>(); + parts.add("table: " + tableName.toLowerCase(Locale.ROOT)); + if (StringUtils.isNotEmpty(comment)) { + parts.add("comment: " + comment); + } + if (table.getSd() != null && table.getSd().getCols() != null) { + for (FieldSchema column : table.getSd().getCols()) { + parts.add(formatColumnForSearch(column)); + } + } + return String.join("; ", parts); + } + + private static String formatColumnsForStorage(Table table) { + if (table.getSd() == null || table.getSd().getCols() == null) { + return ""; + } + List parts = new ArrayList<>(table.getSd().getCols().size()); + for (FieldSchema column : table.getSd().getCols()) { + parts.add(formatColumnForStorage(column)); + } + return String.join("; ", parts); + } + + private static String formatColumnNamesForSearch(Table table) { + if (table.getSd() == null || table.getSd().getCols() == null) { + return ""; + } + List names = new ArrayList<>(table.getSd().getCols().size()); + for (FieldSchema column : table.getSd().getCols()) { + names.add(column.getName().toLowerCase(Locale.ROOT)); + } + return String.join(" ", names); + } + + private static String formatColumnCommentsForSearch(Table table) { + if (table.getSd() == null || table.getSd().getCols() == null) { + return ""; + } + List parts = new ArrayList<>(); + for (FieldSchema column : table.getSd().getCols()) { + if (StringUtils.isNotEmpty(column.getComment())) { + parts.add(column.getComment()); + } + } + return String.join("; ", parts); + } + + private static String formatColumnForStorage(FieldSchema column) { + String base = column.getName() + " " + column.getType(); + if (StringUtils.isNotEmpty(column.getComment())) { + return base + " " + column.getComment(); + } + return base; + } + + private static String formatColumnForSearch(FieldSchema column) { + String name = column.getName().toLowerCase(Locale.ROOT); + if (StringUtils.isNotEmpty(column.getComment())) { + return "column " + name + ": " + column.getComment(); + } + return "column " + name + ":"; + } + + private static String nullToEmpty(String value) { + return value == null ? "" : value; + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/HybridSearch.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/HybridSearch.java new file mode 100644 index 000000000000..33f1f4d8e468 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/HybridSearch.java @@ -0,0 +1,92 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.search; + +import java.util.List; + +import org.apache.commons.lang3.StringUtils; +import org.apache.hive.search.config.SearchConfig; +import org.apache.hive.search.exception.SearchException; +import org.apache.hive.search.mapping.FieldSchema; +import org.apache.hive.search.mapping.IndexMapping; + +public final class HybridSearch { + private HybridSearch() {} + + public record ResolvedHybridQuery(String field, String queryText, Float semanticWeight) {} + + public static ResolvedHybridQuery resolve(SearchArgs.Hybrid args, IndexMapping mapping) + throws SearchException { + String field = requireHybridField(mapping, args.field()); + Float semanticWeight = resolveSemanticWeight(args.matchWeight(), args.semanticWeight()); + return new ResolvedHybridQuery(field, args.queryText(), semanticWeight); + } + + public static SearchInternal.FusionRequest toFusionRequest( + ResolvedHybridQuery hybrid, int defaultSize, SearchConfig searchConfig) + throws SearchException { + float semanticWeight = hybrid.semanticWeight() != null ? + hybrid.semanticWeight() : searchConfig.getHybridSemanticWeight(); + float matchWeight = 1.0f - semanticWeight; + List retrievers = + List.of( + new SearchInternal.RetrieverSpec( + new SearchArgs.Match(hybrid.queryText()), matchWeight, SearchQuery.Mode.MATCH.name()), + new SearchInternal.RetrieverSpec( + new SearchArgs.Semantic(hybrid.queryText(), hybrid.field()), + semanticWeight, + SearchQuery.Mode.SEMANTIC.name()) + ); + return new SearchInternal.FusionRequest(retrievers, defaultSize); + } + + private static Float resolveSemanticWeight(Float matchWeight, Float semanticWeight) { + if (semanticWeight != null) { + return semanticWeight; + } + if (matchWeight != null) { + return 1.0f - matchWeight; + } + return null; + } + + private static String requireHybridField(IndexMapping mapping, String field) + throws SearchException { + if (StringUtils.isNotEmpty(field)) { + validateHybridField(mapping, field); + return field; + } + return mapping + .soleHybridField() + .orElseThrow( + () -> + new SearchException( + "hybrid query requires field when index has " + + mapping.hybridFields().size() + + " hybrid field(s): " + + mapping.hybridFields())); + } + + private static void validateHybridField(IndexMapping mapping, String field) + throws SearchException { + FieldSchema schema = mapping.fieldSchema(field); + if (!(schema instanceof FieldSchema.TextFieldSchema text) || !text.search().hybrid()) { + throw new SearchException("field '" + field + "' is not configured for hybrid search"); + } + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/LuceneSearchBackend.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/LuceneSearchBackend.java new file mode 100644 index 000000000000..f84386cdca96 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/LuceneSearchBackend.java @@ -0,0 +1,108 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.search; + +import java.io.IOException; +import java.util.List; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hive.search.config.SearchConfig; +import org.apache.hive.search.exception.IndexNotReadyException; +import org.apache.hive.search.exception.InitializeException; +import org.apache.hive.search.exception.SearchException; +import org.apache.hive.search.index.IndexSession; +import org.apache.hive.search.metastore.MetastoreTableMapper; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Phase 1 {@link SearchBackend} backed by an in-process Lucene index. */ +public final class LuceneSearchBackend implements SearchBackend { + private static final Logger LOG = LoggerFactory.getLogger(LuceneSearchBackend.class); + private IndexSession session; + private Future future; + private SearchConfig searchConfig; + + @Override + public void initialize(Configuration configuration) + throws InitializeException, IOException { + searchConfig = new SearchConfig(configuration); + session = new IndexSession(configuration); + ExecutorService initThread = Executors.newFixedThreadPool(1); + future = initThread.submit(() -> { + long start = System.currentTimeMillis(); + session.initialize(); + LOG.info("The in-process Lucene search backend is ready for request now, time taken: {}ms", + System.currentTimeMillis() - start); + return null; + }); + initThread.shutdown(); + } + + @Override + public boolean isReady() throws IndexNotReadyException { + if (future == null) { + throw new IndexNotReadyException("The in-process Lucene search backend hasn't been initialized yet"); + } + try { + future.get(searchConfig.getInitReadyTimeoutMs(), TimeUnit.MILLISECONDS); + return true; + } catch (ExecutionException e) { + Throwable cause = e.getCause() != null ? e.getCause() : e; + throw new IndexNotReadyException("Search backend initialization failed", cause); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IndexNotReadyException("Search backend initialization interrupted", e); + } catch (TimeoutException e) { + return false; + } + } + + @Override + public TableSearchResult search(SearchQuery query) + throws SearchException, IOException { + if (!isReady()) { + throw new IndexNotReadyException("Search index is not ready"); + } + List fields = query.returnFields(); + if (query.returnFields().isEmpty()) { + fields = List.of( + MetastoreTableMapper.FIELD_DB, + MetastoreTableMapper.FIELD_TABLE, + MetastoreTableMapper.FIELD_OWNER, + MetastoreTableMapper.FIELD_COMMENT); + } + int limit = query.limit() > 0 ? query.limit() : searchConfig.getDefaultLimit(); + query = SearchQuery.of(query, fields, limit); + try (SearchInternal searcher = session.getSearcher()) { + return searcher.search(query); + } + } + + @Override + public void close() throws Exception { + if (session != null) { + session.close(); + } + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/SearchArgs.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/SearchArgs.java new file mode 100644 index 000000000000..82010de7cb1e --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/SearchArgs.java @@ -0,0 +1,229 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.search; + +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +import org.apache.commons.lang3.StringUtils; +import org.apache.hive.search.exception.SearchException; + +/** Typed search arguments. Map/JSON parsing belongs in {@link #fromBody(Object, SearchQuery.Mode)}. */ +public sealed interface SearchArgs permits SearchArgs.Match, SearchArgs.Semantic, SearchArgs.Hybrid { + + record Match(String queryText) implements SearchArgs { + public Match { + Objects.requireNonNull(queryText, "queryText"); + } + } + + record Semantic(String queryText, String field) implements SearchArgs { + public Semantic { + Objects.requireNonNull(queryText, "queryText"); + } + } + + record Hybrid(String queryText, String field, Float matchWeight, Float semanticWeight) + implements SearchArgs { + public Hybrid { + Objects.requireNonNull(queryText, "queryText"); + } + } + + /** Deserializes an API/Thrift request body into typed search args. */ + static SearchArgs fromBody(Object body, SearchQuery.Mode mode) throws SearchException { + return switch (mode) { + case MATCH -> parseMatchBody(body); + case SEMANTIC -> parseSemanticBody(body); + case HYBRID -> parseHybridBody(body); + }; + } + + /** Serializes typed search args to a flat string map for Thrift. */ + static Map toBody(SearchArgs args) { + return switch (args) { + case Match match -> Map.of("query", match.queryText()); + case Semantic semantic -> toSemanticBody(semantic); + case Hybrid hybrid -> toHybridBody(hybrid); + }; + } + + /** Serializes typed search args to a Thrift-ready query body map. */ + static Map toQueryBody(SearchArgs args) { + return toBody(args); + } + + private static Map toSemanticBody(Semantic semantic) { + Map body = new LinkedHashMap<>(); + body.put("query", semantic.queryText()); + if (StringUtils.isNotEmpty(semantic.field())) { + body.put("field", semantic.field()); + } + return body; + } + + private static Map toHybridBody(Hybrid hybrid) { + Map body = new LinkedHashMap<>(); + body.put("query", hybrid.queryText()); + if (StringUtils.isNotEmpty(hybrid.field())) { + body.put("field", hybrid.field()); + } + if (hybrid.matchWeight() != null) { + body.put("match_weight", hybrid.matchWeight().toString()); + } + if (hybrid.semanticWeight() != null) { + body.put("semantic_weight", hybrid.semanticWeight().toString()); + } + return body; + } + + private static Match parseMatchBody(Object body) throws SearchException { + if (body instanceof String queryText) { + return new Match(queryText); + } + if (body instanceof Map matchMap && matchMap.containsKey("query")) { + return new Match(matchMap.get("query").toString()); + } + throw new SearchException("match query must be a string or {query: text}"); + } + + private static Semantic parseSemanticBody(Object body) throws SearchException { + if (body instanceof String queryText) { + return new Semantic(queryText, null); + } + if (!(body instanceof Map semanticMap)) { + throw new SearchException("semantic query must be a string or object"); + } + if (semanticMap.containsKey("field") && semanticMap.containsKey("query")) { + return new Semantic( + semanticMap.get("query").toString(), + semanticMap.get("field").toString()); + } + if (semanticMap.containsKey("query") && !semanticMap.containsKey("field")) { + return new Semantic(semanticMap.get("query").toString(), null); + } + var fieldEntries = semanticMap.entrySet().stream() + .filter(entry -> entry.getValue() != null) + .toList(); + if (fieldEntries.size() == 1) { + Map.Entry entry = fieldEntries.getFirst(); + return new Semantic(entry.getValue().toString(), entry.getKey().toString()); + } + throw new SearchException("semantic query must be a string, {field:, query:}, or {: text}"); + } + + private static Hybrid parseHybridBody(Object body) throws SearchException { + if (body instanceof String queryText) { + return new Hybrid(queryText, null, null, null); + } + if (!(body instanceof Map hybridMap)) { + throw new SearchException("hybrid query must be a string or object"); + } + @SuppressWarnings("unchecked") + Map hybridBody = (Map) hybridMap; + Float matchWeight = parseHybridWeight(hybridBody, "match_weight", "match"); + Float semanticWeight = parseHybridWeight(hybridBody, "semantic_weight", "semantic"); + validateHybridWeights(matchWeight, semanticWeight); + + if (hybridBody.containsKey("field") && hybridBody.containsKey("query")) { + return new Hybrid( + hybridBody.get("query").toString(), + hybridBody.get("field").toString(), + matchWeight, + semanticWeight); + } + if (hybridBody.containsKey("query")) { + return new Hybrid( + hybridBody.get("query").toString(), + null, + matchWeight, + semanticWeight); + } + List> fieldEntries = hybridBody.entrySet().stream() + .filter(entry -> !isHybridOption(entry.getKey())) + .filter(entry -> entry.getValue() != null) + .toList(); + if (fieldEntries.size() == 1) { + Map.Entry entry = fieldEntries.getFirst(); + return new Hybrid( + entry.getValue().toString(), + entry.getKey(), + matchWeight, + semanticWeight); + } + throw new SearchException( + "hybrid query must be a string, {field:, query:}, {query:}, or {: text}"); + } + + private static Float parseHybridWeight( + Map hybridBody, String flatKey, String nestedKey) + throws SearchException { + Object flatValue = hybridBody.get(flatKey); + if (flatValue != null) { + return parseWeightValue(flatValue, flatKey); + } + if (hybridBody.get("weights") instanceof Map weights) { + Object nestedValue = weights.get(nestedKey); + if (nestedValue != null) { + return parseWeightValue(nestedValue, nestedKey); + } + } + return null; + } + + private static Float parseWeightValue(Object value, String key) throws SearchException { + if (value instanceof Number number) { + return number.floatValue(); + } + if (value instanceof String text) { + try { + return Float.parseFloat(text); + } catch (NumberFormatException e) { + throw new SearchException("Invalid hybrid weight for '" + key + "': " + text); + } + } + throw new SearchException("Invalid hybrid weight for '" + key + "': " + value); + } + + private static void validateHybridWeights(Float matchWeight, Float semanticWeight) + throws SearchException { + if (!isWeightValid(matchWeight)) { + throw new SearchException( + "Invalid hybrid lexical weight, it must be in (0, 1), but got " + matchWeight); + } + if (!isWeightValid(semanticWeight)) { + throw new SearchException( + "Invalid hybrid semantic weight, it must be in (0, 1), but got " + semanticWeight); + } + if (matchWeight != null && semanticWeight != null && (matchWeight + semanticWeight) > 1.0f) { + throw new SearchException("Invalid hybrid weights, sum of weights exceeds 1.0f"); + } + } + + private static boolean isWeightValid(Float weight) { + return weight == null || (weight > 0.0f && weight < 1.0f); + } + + private static boolean isHybridOption(String key) { + return "match_weight".equals(key) + || "semantic_weight".equals(key) + || "weights".equals(key); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/SearchBackend.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/SearchBackend.java new file mode 100644 index 000000000000..d85e8788f1e2 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/SearchBackend.java @@ -0,0 +1,44 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.search; + +import java.io.IOException; + +import org.apache.hadoop.conf.Configuration; + +import org.apache.hive.search.exception.IndexNotReadyException; +import org.apache.hive.search.exception.InitializeException; +import org.apache.hive.search.exception.SearchException; + +/** + * Pluggable search storage and query engine. Phase 1 uses embedded Lucene; Phase 3 may add + * OpenSearch or Elasticsearch without changing the mapping contract or ingest path. + */ +public interface SearchBackend extends AutoCloseable { + + void initialize(Configuration configuration) throws InitializeException, IOException; + + /** Whether this instance can serve search (index opened and bootstrap completed). */ + boolean isReady() throws IndexNotReadyException; + + TableSearchResult search(SearchQuery query) + throws SearchException, InitializeException, IOException; + + @Override + void close() throws Exception; +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/SearchInternal.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/SearchInternal.java new file mode 100644 index 000000000000..cc3355727a68 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/SearchInternal.java @@ -0,0 +1,288 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.search; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Objects; + +import org.apache.commons.lang3.StringUtils; +import org.apache.hadoop.hive.common.TableName; +import org.apache.hive.search.config.SearchConfig; +import org.apache.hive.search.exception.IndexException; +import org.apache.hive.search.exception.SearchException; +import org.apache.hive.search.mapping.FieldSchema; +import org.apache.hive.search.mapping.IndexMapping; +import org.apache.hive.search.mapping.TableDocument; +import org.apache.hive.search.metastore.MetastoreTableMapper; +import org.apache.hive.search.index.IndexManager; +import org.apache.hive.search.inference.EmbedModel; +import org.apache.hive.search.inference.EmbedModelRegistry; +import org.apache.lucene.analysis.Analyzer; +import org.apache.lucene.analysis.TokenStream; +import org.apache.lucene.analysis.tokenattributes.CharTermAttribute; +import org.apache.lucene.document.Document; +import org.apache.lucene.index.IndexableField; +import org.apache.lucene.index.Term; +import org.apache.lucene.search.BayesianScoreEstimator; +import org.apache.lucene.search.BayesianScoreQuery; +import org.apache.lucene.search.BooleanClause; +import org.apache.lucene.search.BooleanQuery; +import org.apache.lucene.search.BoostQuery; +import org.apache.lucene.search.ConstantScoreQuery; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.LogOddsFusionQuery; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.SearcherManager; +import org.apache.lucene.search.TermQuery; +import org.apache.lucene.search.PrefixQuery; +import org.apache.lucene.search.TopDocs; + +public final class SearchInternal implements AutoCloseable { + private final EmbedModelRegistry modelRegistry; + private final IndexSearcher searcher; + private final SearcherManager searcherManager; + private final IndexMapping mapping; + private final SearchConfig searchConfig; + private final BayesianScoreEstimator.Parameters parameters; + private final long indexedNid; + + public SearchInternal(SearcherManager manager, + IndexManager indexManager, + EmbedModelRegistry registry, + SearchConfig searchConfig, + BayesianScoreEstimator.Parameters parameters) throws IOException { + this.searcherManager = manager; + this.indexedNid = indexManager.getIndexedNid(); + this.searcher = manager.acquire(); + this.mapping = indexManager.mapping(); + this.searchConfig = searchConfig; + this.parameters = parameters; + this.modelRegistry = Objects.requireNonNull(registry, "Model registry"); + } + + public record RetrieverSpec(SearchArgs query, float weight, String name) { + } + + public record FusionRequest(List retrievers, int size) {} + + public TableSearchResult search(SearchQuery request) + throws SearchException, IOException { + int size = request.limit() > 0 ? request.limit() : searchConfig.getDefaultLimit(); + Query query = compileRootQuery(request, size); + query = applyScopeFilter(query, request.catalogName(), request.databaseName()); + TopDocs topDocs = searcher.search(query, size); + List hits = new ArrayList<>(); + for (ScoreDoc scoreDoc : topDocs.scoreDocs) { + hits.add(readHit(scoreDoc, request.returnFields())); + } + return new TableSearchResult(hits, topDocs.totalHits.value(), indexedNid); + } + + private Query compileRootQuery(SearchQuery searchQuery, int size) + throws SearchException, IOException { + return switch (searchQuery.args()) { + case SearchArgs.Match match -> compileMatchQuery(match.queryText()); + case SearchArgs.Semantic semantic -> + compileSemanticQuery(SemanticSearch.resolve(semantic, mapping), size); + case SearchArgs.Hybrid hybrid -> { + HybridSearch.ResolvedHybridQuery resolved = HybridSearch.resolve(hybrid, mapping); + yield compileFusionQuery( + HybridSearch.toFusionRequest(resolved, size, searchConfig), size); + } + }; + } + + private Query compileFusionQuery(FusionRequest fusion, int size) + throws SearchException, IOException { + int i = 0; + float[] weights = new float[fusion.retrievers.size()]; + List queries = new ArrayList<>(); + for (RetrieverSpec retriever : fusion.retrievers()) { + Query internalQuery = compileQuery(retriever.query(), size); + if (internalQuery instanceof KnnFloatVectorQuery) { + queries.add(internalQuery); + } else { + queries.add(new BayesianScoreQuery(internalQuery, parameters.alpha(), + parameters.beta(), parameters.baseRate())); + } + weights[i++] = retriever.weight(); + } + return new LogOddsFusionQuery(queries, searchConfig.getFusionPrior(), weights); + } + + private Query applyScopeFilter(Query query, String catalogName, String databaseName) { + Query filter = compileScopeFilter(catalogName, databaseName); + if (filter == null) { + return query; + } + return new BooleanQuery.Builder() + .add(query, BooleanClause.Occur.MUST) + .add(filter, BooleanClause.Occur.FILTER) + .build(); + } + + private Query compileScopeFilter(String catalogName, String databaseName) { + boolean hasCatalog = StringUtils.isNotEmpty(catalogName); + boolean hasDatabase = StringUtils.isNotEmpty(databaseName); + if (!hasCatalog && !hasDatabase) { + return null; + } + if (hasCatalog && hasDatabase) { + return new PrefixQuery(new Term("_id" + TableDocument.FILTER_SUFFIX, + catalogName + "." + databaseName + ".")); + } + if (hasCatalog) { + return new PrefixQuery(new Term("_id" + TableDocument.FILTER_SUFFIX, catalogName + ".")); + } + return new TermQuery(new Term( + MetastoreTableMapper.FIELD_DB + TableDocument.FILTER_SUFFIX, databaseName)); + } + + private Query compileQuery(SearchArgs args, int knnK) + throws SearchException, IOException { + return switch (args) { + case SearchArgs.Match match -> compileMatchQuery(match.queryText()); + case SearchArgs.Semantic semantic -> + compileSemanticQuery(SemanticSearch.resolve(semantic, mapping), knnK); + case SearchArgs.Hybrid hybrid -> + throw new SearchException("Nested hybrid queries are not supported"); + }; + } + + private Query compileMatchQuery(String queryText) throws SearchException, IOException { + BooleanQuery.Builder builder = new BooleanQuery.Builder(); + boolean added = false; + for (MetastoreTableMapper.KeywordSearchField keywordField : MetastoreTableMapper.KEYWORD_SEARCH_FIELDS) { + String field = keywordField.field(); + float boost = keywordField.boost(); + FieldSchema schema = mapping.fieldSchema(field); + if (!(schema instanceof FieldSchema.TextFieldSchema text)) { + continue; + } + if (text.filter()) { + builder.add( + boostKeywordQuery(compileFilterKeywordQuery(field, queryText), boost), + BooleanClause.Occur.SHOULD); + if (text.search().lexical()) { + builder.add( + boostKeywordQuery(compileMatchQuery(field, queryText), boost), + BooleanClause.Occur.SHOULD); + } + added = true; + continue; + } + if (!text.search().lexical()) { + continue; + } + builder.add( + boostKeywordQuery(compileMatchQuery(field, queryText), boost), + BooleanClause.Occur.SHOULD); + added = true; + } + if (!added) { + throw new SearchException("No lexically searchable table fields are configured"); + } + return builder.build(); + } + + private static Query boostKeywordQuery(Query query, float boost) { + return new BoostQuery(new ConstantScoreQuery(query), boost); + } + + private Query compileFilterKeywordQuery(String field, String queryText) { + String normalized = queryText.trim().toLowerCase(Locale.ROOT); + if (normalized.isEmpty()) { + return new BooleanQuery.Builder().build(); + } + return new TermQuery(new Term(field + TableDocument.FILTER_SUFFIX, normalized)); + } + + private Query compileMatchQuery(String field, String queryText) + throws SearchException, IOException { + + FieldSchema schema = mapping.fieldSchema(field); + if (!(schema instanceof FieldSchema.TextFieldSchema text) || !text.search().lexical()) { + throw new SearchException("field '" + field + "' is not lexically searchable"); + } + + Analyzer analyzer = mapping.analyzer(); + BooleanQuery.Builder builder = new BooleanQuery.Builder(); + try (TokenStream stream = analyzer.tokenStream(field, queryText)) { + CharTermAttribute termAttr = stream.addAttribute(CharTermAttribute.class); + stream.reset(); + while (stream.incrementToken()) { + builder.add(new TermQuery(new Term(field, termAttr.toString())), BooleanClause.Occur.SHOULD); + } + stream.end(); + } + return builder.build(); + } + + private org.apache.lucene.search.Query compileSemanticQuery( + SemanticSearch.ResolvedSemanticQuery semantic, int knnK) throws SearchException { + FieldSchema schema = mapping.fieldSchema(semantic.field()); + if (!(schema instanceof FieldSchema.TextFieldSchema text) || !text.search().semantic()) { + throw new SearchException("field '" + semantic.field() + "' is not semantically searchable"); + } + float[] embedding; + try { + embedding = modelRegistry.get(text.search().semanticModel()) + .embed(EmbedModel.TaskType.QUERY, semantic.queryText()); + } catch (IndexException e) { + throw new SearchException( + "Failed to encode semantic query for field '" + semantic.field() + "'", e); + } + if (embedding == null) { + throw new SearchException( + "embedding model '" + text.search().semanticModel() + "' returned null vector"); + } + return new KnnFloatVectorQuery(semantic.field(), embedding, knnK); + } + + private TableSearchHit readHit(ScoreDoc scoreDoc, List fields) + throws IOException { + Document stored = searcher.storedFields().document(scoreDoc.doc); + Map fieldHits = new LinkedHashMap<>(); + List requested = fields.isEmpty() ? + mapping.fields().keySet().stream().toList() : fields; + for (String field : requested) { + IndexableField[] values = stored.getFields(field); + if (values != null && values.length > 0) { + fieldHits.put(field, values[0].stringValue()); + } + } + TableName tableName = null; + IndexableField[] idValues = stored.getFields("_id"); + if (idValues != null && idValues.length > 0) { + tableName = TableName.fromString(idValues[0].stringValue(), "", "default"); + } + return new TableSearchHit(tableName, scoreDoc.score, fieldHits); + } + + @Override + public void close() throws IOException { + searcherManager.release(searcher); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/SearchQuery.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/SearchQuery.java new file mode 100644 index 000000000000..976d58938a09 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/SearchQuery.java @@ -0,0 +1,117 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.search; + +import java.util.List; +import java.util.Map; + +import org.apache.commons.lang3.StringUtils; +import org.apache.hive.search.exception.SearchException; + +public record SearchQuery( + SearchArgs args, + String catalogName, + String databaseName, + int limit, + List returnFields) { + + /** Public search mode exposed on the Metastore Thrift API. */ + public enum Mode { + MATCH, + SEMANTIC, + HYBRID + } + + public SearchQuery { + returnFields = returnFields == null ? List.of() : List.copyOf(returnFields); + } + + public Mode mode() { + return switch (args) { + case SearchArgs.Match m -> Mode.MATCH; + case SearchArgs.Semantic s -> Mode.SEMANTIC; + case SearchArgs.Hybrid h -> Mode.HYBRID; + }; + } + + private static void validate(SearchArgs args, int limit) throws SearchException { + String queryText = switch (args) { + case SearchArgs.Match m -> m.queryText(); + case SearchArgs.Semantic s -> s.queryText(); + case SearchArgs.Hybrid h -> h.queryText(); + }; + validate(queryText, limit); + } + + private static void validate(String queryText, int limit) throws SearchException { + if (StringUtils.isEmpty(queryText)) { + throw new SearchException("queryText is required"); + } + if (limit < 0) { + throw new SearchException("limit must be non-negative"); + } + } + + public static SearchQuery fromQueryBody( + Map queryBody, + Mode mode, + String catalogName, + String databaseName, + int limit, + List returnFields) + throws SearchException { + if (queryBody == null || queryBody.isEmpty()) { + throw new SearchException("missing query body for mode " + mode); + } + SearchArgs args = SearchArgs.fromBody(queryBody, mode); + validate(args, limit); + return new SearchQuery(args, catalogName, databaseName, limit, returnFields); + } + + public static SearchQuery of(String queryText) throws SearchException { + validate(queryText, 0); + return new SearchQuery(new SearchArgs.Hybrid(queryText, null, null, null), null, null, 0, List.of()); + } + + public static SearchQuery of(String queryText, Mode mode, int limit) throws SearchException { + validate(queryText, limit); + SearchArgs args = switch (mode) { + case MATCH -> new SearchArgs.Match(queryText); + case SEMANTIC -> new SearchArgs.Semantic(queryText, null); + case HYBRID -> new SearchArgs.Hybrid(queryText, null, null, null); + }; + return new SearchQuery(args, null, null, limit, List.of()); + } + + public static SearchQuery of(String keyWord, String catalogName, String databaseName) + throws SearchException { + validate(keyWord, 0); + return new SearchQuery( + new SearchArgs.Match(keyWord), catalogName, databaseName, 0, List.of()); + } + + public static SearchQuery of(SearchQuery query, List returnFields, int limit) { + return new SearchQuery( + query.args(), query.catalogName(), query.databaseName(), limit, returnFields); + } + + /** Serializes this query to a flat Thrift-ready query body map. */ + public Map toQueryBody() { + return SearchArgs.toQueryBody(args); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/SemanticSearch.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/SemanticSearch.java new file mode 100644 index 000000000000..8645119a755c --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/SemanticSearch.java @@ -0,0 +1,66 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.search; + +import java.util.List; + +import org.apache.commons.lang3.StringUtils; +import org.apache.hive.search.exception.SearchException; +import org.apache.hive.search.mapping.FieldSchema; +import org.apache.hive.search.mapping.IndexMapping; + +public final class SemanticSearch { + private SemanticSearch() {} + + public record ResolvedSemanticQuery(String field, String queryText) {} + + public static ResolvedSemanticQuery resolve(SearchArgs.Semantic args, IndexMapping mapping) + throws SearchException { + String field = requireSemanticField(mapping, args.field()); + return new ResolvedSemanticQuery(field, args.queryText()); + } + + private static String requireSemanticField(IndexMapping mapping, String field) + throws SearchException { + if (StringUtils.isNotEmpty(field)) { + validateSemanticField(mapping, field); + return field; + } + List semanticFields = mapping.fields().entrySet().stream() + .filter(entry -> entry.getValue() instanceof FieldSchema.TextFieldSchema text + && text.search().semantic()) + .map(java.util.Map.Entry::getKey) + .toList(); + if (semanticFields.size() == 1) { + return semanticFields.getFirst(); + } + throw new SearchException( + "semantic query requires field when index has " + + semanticFields.size() + + " semantic field(s): " + + semanticFields); + } + + private static void validateSemanticField(IndexMapping mapping, String field) + throws SearchException { + FieldSchema schema = mapping.fieldSchema(field); + if (!(schema instanceof FieldSchema.TextFieldSchema text) || !text.search().semantic()) { + throw new SearchException("field '" + field + "' is not semantically searchable"); + } + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/TableSearchHit.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/TableSearchHit.java new file mode 100644 index 000000000000..4ae0d07269fa --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/TableSearchHit.java @@ -0,0 +1,36 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.search; + +import java.util.Map; + +import org.apache.hadoop.hive.common.TableName; + +/** One ranked discovery hit. Schema detail requires a canonical Metastore read. */ +public record TableSearchHit( + TableName table, float score, Map fields) { + + public TableSearchHit { + fields = fields == null ? Map.of() : Map.copyOf(fields); + } + + @Override + public String toString() { + return "TableSearchHit{" + "table=" + table + ", score=" + score + '}'; + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/TableSearchResult.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/TableSearchResult.java new file mode 100644 index 000000000000..3c6b1f6d5bfb --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/TableSearchResult.java @@ -0,0 +1,26 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.search; + +import java.util.List; + +public record TableSearchResult(List hits, long total, long indexedNid) { + public TableSearchResult { + hits = hits == null ? List.of() : List.copyOf(hits); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hadoop/hive/metastore/handler/TestSearchProvider.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hadoop/hive/metastore/handler/TestSearchProvider.java new file mode 100644 index 000000000000..9d8a185b08e3 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hadoop/hive/metastore/handler/TestSearchProvider.java @@ -0,0 +1,116 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hadoop.hive.metastore.handler; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.common.TableName; +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hive.search.search.SearchBackend; +import org.apache.hive.search.search.SearchQuery; +import org.apache.hive.search.search.TableSearchHit; +import org.apache.hive.search.search.TableSearchResult; +import org.junit.After; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.util.List; +import java.util.Map; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +@Category(MetastoreUnitTest.class) +public class TestSearchProvider { + + @After + public void tearDown() throws Exception { + SearchProvider.reset(); + } + + @Test + public void createInitializesBackend() throws Exception { + StubSearchBackend backend = new StubSearchBackend(); + SearchProvider provider = + SearchProvider.install(new Configuration(false), backend); + + assertTrue(backend.initialized); + assertTrue(provider.isReady()); + assertSame(backend, provider.backend()); + provider.close(); + } + + @Test + public void installExposesSingleton() throws Exception { + Configuration conf = new Configuration(false); + SearchProvider installed = + SearchProvider.install(conf, new StubSearchBackend()); + assertSame(installed, SearchProvider.get()); + } + + @Test + public void getFailsWhenNotInstalled() { + assertThrows(IllegalStateException.class, SearchProvider::get); + } + + @Test + public void searchDelegatesToBackend() throws Exception { + StubSearchBackend backend = new StubSearchBackend(); + SearchProvider provider = + SearchProvider.install(new Configuration(false), backend); + SearchQuery query = SearchQuery.of("sales", SearchQuery.Mode.MATCH, 5); + + TableSearchResult result = provider.search(query); + + assertEquals(1, result.total()); + assertEquals("t", result.hits().getFirst().table().getTable()); + provider.close(); + } + + private static final class StubSearchBackend implements SearchBackend { + private boolean initialized; + private boolean closed; + + @Override + public void initialize(Configuration configuration) { + initialized = true; + } + + @Override + public boolean isReady() { + return initialized; + } + + @Override + public TableSearchResult search(SearchQuery query) { + return new TableSearchResult( + List.of(new TableSearchHit( + TableName.fromString("default.t", "hive", "default"), + 1.0f, + Map.of("table_name", "t"))), + 1, 1); + } + + @Override + public void close() { + closed = true; + initialized = false; + } + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/config/TestIndexConfig.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/config/TestIndexConfig.java new file mode 100644 index 000000000000..6cc4332aefd1 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/config/TestIndexConfig.java @@ -0,0 +1,49 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.config; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.time.Duration; + +import static org.junit.Assert.assertEquals; + +@Category(MetastoreUnitTest.class) +public class TestIndexConfig { + + @Test + public void usesDefaultsWhenUnset() { + IndexConfig config = new IndexConfig(new Configuration(false)); + assertEquals(IndexConfig.INDEX_NAME_DEFAULT, config.indexName()); + assertEquals(IndexConfig.INDEX_RAM_SIZE_DEFAULT, config.getWriteBufferSize()); + assertEquals(Duration.ofMillis(IndexConfig.FLUSH_INTERVAL_MS_DEFAULT), config.getFlushInterval()); + } + + @Test + public void readsConfiguredValues() { + Configuration conf = new Configuration(false); + conf.set(IndexConfig.INDEX_NAME, "custom_index"); + conf.setInt(IndexConfig.BOOTSTRAP_BATCH_SIZE, 500); + IndexConfig config = new IndexConfig(conf); + assertEquals("custom_index", config.indexName()); + assertEquals(500, config.getBootstrapBatchSize()); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/config/TestIndexStateConfig.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/config/TestIndexStateConfig.java new file mode 100644 index 000000000000..f4657f025732 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/config/TestIndexStateConfig.java @@ -0,0 +1,54 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.config; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hive.search.exception.IndexIOException; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +@Category(MetastoreUnitTest.class) +public class TestIndexStateConfig { + + @Test + public void detectsRemoteAndMemoryFlags() { + Configuration conf = new Configuration(false); + conf.set(IndexStoreConfig.REMOTE_URI, "file:///tmp/backup"); + conf.setBoolean(IndexStoreConfig.MEMORY, true); + IndexStoreConfig config = new IndexStoreConfig(conf, "hive_tables"); + assertTrue(config.hasRemote()); + assertTrue(config.isDistributed()); + assertTrue(config.useMemory()); + } + + @Test + public void validateRemoteUriRejectsInvalidValue() { + assertThrows(IndexIOException.class, () -> IndexStoreConfig.validateRemoteUri("://bad")); + } + + @Test + public void localPathDefaultsToWorkdir() { + IndexStoreConfig config = new IndexStoreConfig(new Configuration(false), "hive_tables"); + assertFalse(config.getLocalPath().toString().isEmpty()); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/config/TestInferenceConfig.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/config/TestInferenceConfig.java new file mode 100644 index 000000000000..5846a56ae2fa --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/config/TestInferenceConfig.java @@ -0,0 +1,68 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.config; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +@Category(MetastoreUnitTest.class) +public class TestInferenceConfig { + + @Test + public void embeddingThreadsDefaultsToAtLeastOne() { + Configuration conf = new Configuration(false); + InferenceConfig config = new InferenceConfig(conf); + assertTrue(config.getEmbeddingThreads() >= 1); + } + + @Test + public void embeddingThreadsHonorsConfiguration() { + Configuration conf = new Configuration(false); + conf.setInt(InferenceConfig.EMBEDDING_THREADS, 8); + InferenceConfig config = new InferenceConfig(conf); + assertEquals(8, config.getEmbeddingThreads()); + } + + @Test + public void embeddingThreadsRejectsNonPositiveValues() { + Configuration conf = new Configuration(false); + conf.setInt(InferenceConfig.EMBEDDING_THREADS, 0); + InferenceConfig config = new InferenceConfig(conf); + assertEquals(1, config.getEmbeddingThreads()); + } + + @Test + public void embeddingMaxSeqLengthDefaultsTo512() { + Configuration conf = new Configuration(false); + InferenceConfig config = new InferenceConfig(conf); + assertEquals(512, config.getEmbeddingMaxSeqLength()); + } + + @Test + public void embeddingMaxSeqLengthHonorsConfiguration() { + Configuration conf = new Configuration(false); + conf.setInt(InferenceConfig.EMBEDDING_MAX_SEQ_LENGTH, 256); + InferenceConfig config = new InferenceConfig(conf); + assertEquals(256, config.getEmbeddingMaxSeqLength()); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/config/TestSearchConfig.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/config/TestSearchConfig.java new file mode 100644 index 000000000000..23d8d067f187 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/config/TestSearchConfig.java @@ -0,0 +1,51 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.config; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hive.search.exception.SearchException; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.time.Duration; + +import static org.junit.Assert.assertEquals; + +@Category(MetastoreUnitTest.class) +public class TestSearchConfig { + + @Test + public void usesDefaultsWhenUnset() throws SearchException { + SearchConfig config = new SearchConfig(new Configuration(false)); + assertEquals(Duration.ofSeconds(SearchConfig.REFRESH_INTERVAL_SECONDS_DEFAULT), + config.getRefreshInterval()); + assertEquals(SearchConfig.DEFAULT_LIMIT_DEFAULT, config.getDefaultLimit()); + assertEquals(SearchConfig.HYBRID_SEMANTIC_WEIGHT_DEFAULT, config.getHybridSemanticWeight(), 0.001f); + } + + @Test + public void readsConfiguredValues() throws SearchException { + Configuration conf = new Configuration(false); + conf.setInt(SearchConfig.DEFAULT_LIMIT, 50); + conf.setFloat(SearchConfig.HYBRID_SEMANTIC_WEIGHT, 0.2f); + SearchConfig config = new SearchConfig(conf); + assertEquals(50, config.getDefaultLimit()); + assertEquals(0.8f, config.getHybridMatchWeight(), 0.001f); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/index/TestIndexManagerFailureRecovery.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/index/TestIndexManagerFailureRecovery.java new file mode 100644 index 000000000000..eeafab99f177 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/index/TestIndexManagerFailureRecovery.java @@ -0,0 +1,54 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.index; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hive.search.config.IndexConfig; +import org.apache.hive.search.config.IndexStoreConfig; +import org.apache.hive.search.config.InferenceConfig; +import org.apache.hive.search.exception.IndexNotHealthyException; +import org.apache.hive.search.mapping.IndexMapping; +import org.apache.hive.search.metastore.MetastoreSchemas; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import static org.junit.Assert.assertThrows; + +@Category(MetastoreUnitTest.class) +public class TestIndexManagerFailureRecovery { + + @Test + public void unhealthyStatusBlocksSearchUntilRecovered() throws Exception { + Configuration conf = new Configuration(false); + conf.setBoolean(IndexStoreConfig.MEMORY, true); + conf.set(IndexConfig.INDEX_NAME, "test_index"); + conf.set(InferenceConfig.MODEL_NAME, "stub-model"); + IndexMapping mapping = MetastoreSchemas.defaultHiveTablesMapping("test_index", "stub-model", conf); + + try (IndexManager indexManager = IndexManager.open(mapping, conf)) { + IndexNotHealthyException failure = + new IndexNotHealthyException("notification apply failed"); + indexManager.notifyIndexState(false, failure); + assertThrows(IndexNotHealthyException.class, indexManager::checkIndexState); + + indexManager.notifyIndexState(true); + indexManager.checkIndexState(); + } + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/index/TestIndexerEmbeddingCache.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/index/TestIndexerEmbeddingCache.java new file mode 100644 index 000000000000..7b10d921e2d0 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/index/TestIndexerEmbeddingCache.java @@ -0,0 +1,134 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.index; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hadoop.hive.metastore.api.Table; +import org.apache.hive.search.config.IndexConfig; +import org.apache.hive.search.config.IndexStoreConfig; +import org.apache.hive.search.config.InferenceConfig; +import org.apache.hive.search.inference.EmbedModelRegistry; +import org.apache.hive.search.inference.EmbeddingCache; +import org.apache.hive.search.mapping.IndexMapping; +import org.apache.hive.search.mapping.TableDocument; +import org.apache.hive.search.metastore.MetastoreSchemas; +import org.apache.hive.search.metastore.MetastoreTableMapper; +import org.apache.hive.search.search.InMemorySearchFixture; +import org.apache.hive.search.testutil.StubEmbedModel; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.util.List; +import java.util.Map; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +@Category(MetastoreUnitTest.class) +public class TestIndexerEmbeddingCache { + + @Test + public void dedupesIdenticalSearchTextAcrossTables() throws Exception { + try (IndexerFixture fixture = IndexerFixture.create()) { + TableDocument salesOrders = document(fixture.mapping(), + InMemorySearchFixture.table("hive", "sales", "orders", "shared sales text")); + TableDocument inventoryOrders = document(fixture.mapping(), + InMemorySearchFixture.table("hive", "inventory", "orders", "shared sales text")); + + fixture.indexer().embedDocuments(List.of(salesOrders)); + int callsAfterFirst = fixture.model().encodeBatchCalls(); + assertTrue(callsAfterFirst >= 1); + + fixture.indexer().embedDocuments(List.of(inventoryOrders)); + assertEquals(callsAfterFirst, fixture.model().encodeBatchCalls()); + assertEquals(1, fixture.indexer().embeddingCache().hits()); + } + } + + @Test + public void skipsInferenceWhenSearchTextUnchangedOnUpdate() throws Exception { + try (IndexerFixture fixture = IndexerFixture.create()) { + Table table = InMemorySearchFixture.table("hive", "sales", "orders", "daily sales orders"); + TableDocument original = document(fixture.mapping(), table); + fixture.indexer().embedDocuments(List.of(original)); + int callsAfterInitial = fixture.model().encodeBatchCalls(); + + TableDocument updated = document(fixture.mapping(), table); + fixture.indexer().embedDocuments(List.of(updated)); + assertEquals(callsAfterInitial, fixture.model().encodeBatchCalls()); + assertEquals(1, fixture.indexer().embeddingCache().hits()); + } + } + + private static TableDocument document(IndexMapping mapping, Table table) { + return MetastoreTableMapper.fromTable(table, mapping); + } + + private static final class IndexerFixture implements AutoCloseable { + private final IndexManager indexManager; + private final Indexer indexer; + private final IndexMapping mapping; + private final StubEmbedModel model; + + private IndexerFixture( + IndexManager indexManager, Indexer indexer, IndexMapping mapping, StubEmbedModel model) { + this.indexManager = indexManager; + this.indexer = indexer; + this.mapping = mapping; + this.model = model; + } + + static IndexerFixture create() throws Exception { + Configuration conf = new Configuration(false); + conf.setBoolean(IndexStoreConfig.MEMORY, true); + conf.set(IndexConfig.INDEX_NAME, "test_index"); + conf.set(InferenceConfig.MODEL_NAME, InMemorySearchFixture.MODEL_NAME); + conf.setBoolean(InferenceConfig.EMBEDDING_CACHE_ENABLED, true); + conf.setInt(InferenceConfig.EMBEDDING_CACHE_MAX_ENTRIES, 1000); + + IndexMapping mapping = MetastoreSchemas.defaultHiveTablesMapping( + "test_index", InMemorySearchFixture.MODEL_NAME, conf); + StubEmbedModel model = new StubEmbedModel(InMemorySearchFixture.MODEL_NAME); + EmbedModelRegistry registry = + new EmbedModelRegistry(Map.of(model.name(), model), EmbeddingCache.create(conf)); + IndexManager indexManager = IndexManager.open(mapping, conf); + Indexer indexer = new Indexer(indexManager, registry); + indexer.initialize(); + return new IndexerFixture(indexManager, indexer, mapping, model); + } + + IndexMapping mapping() { + return mapping; + } + + Indexer indexer() { + return indexer; + } + + StubEmbedModel model() { + return model; + } + + @Override + public void close() throws Exception { + indexer.close(); + indexManager.close(); + } + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/index/TestIndexerFlushCommit.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/index/TestIndexerFlushCommit.java new file mode 100644 index 000000000000..02fec26658e9 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/index/TestIndexerFlushCommit.java @@ -0,0 +1,125 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.index; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hive.search.config.IndexConfig; +import org.apache.hive.search.config.IndexStoreConfig; +import org.apache.hive.search.config.InferenceConfig; +import org.apache.hive.search.inference.EmbedModelRegistry; +import org.apache.hive.search.mapping.IndexMapping; +import org.apache.hive.search.mapping.TableDocument; +import org.apache.hive.search.metastore.MetastoreSchemas; +import org.apache.hive.search.metastore.MetastoreTableMapper; +import org.apache.hive.search.search.InMemorySearchFixture; +import org.apache.hive.search.testutil.StubEmbedModel; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.io.IOException; +import java.util.List; +import java.util.Map; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +@Category(MetastoreUnitTest.class) +public class TestIndexerFlushCommit { + + @Test + public void commitsAfterConfiguredFlushCount() throws Exception { + Configuration conf = baseConf(); + conf.setInt(IndexConfig.INDEX_RAM_SIZE, 1024); + conf.setInt(IndexConfig.COMMIT_FLUSHES, 2); + + try (IndexerFixture fixture = IndexerFixture.create(conf)) { + fixture.indexer.addDocuments(List.of(tableDoc(fixture.mapping(), "hive.sales.orders", "orders"))); + assertFalse(fixture.indexer.flush(1L, false)); + assertEquals(1, fixture.indexer.flushesSinceCommit()); + + fixture.indexer.addDocuments(List.of(tableDoc(fixture.mapping(), "hive.sales.customers", "customers"))); + assertTrue(fixture.indexer.flush(2L, false)); + assertEquals(0, fixture.indexer.flushesSinceCommit()); + assertEquals(2L, fixture.indexManager.readLocalManifest().orElseThrow().lastEventId()); + } + } + + @Test + public void forceCommitIgnoresFlushThreshold() throws Exception { + Configuration conf = baseConf(); + conf.setInt(IndexConfig.COMMIT_FLUSHES, 5); + + try (IndexerFixture fixture = IndexerFixture.create(conf)) { + fixture.indexer.addDocuments(List.of(tableDoc(fixture.mapping(), "hive.sales.orders", "orders"))); + assertTrue(fixture.indexer.flush(9L, true)); + assertEquals(0, fixture.indexer.flushesSinceCommit()); + assertEquals(9L, fixture.indexManager.readLocalManifest().orElseThrow().lastEventId()); + } + } + + private static Configuration baseConf() { + Configuration conf = new Configuration(false); + conf.setBoolean(IndexStoreConfig.MEMORY, true); + conf.set(IndexConfig.INDEX_NAME, "test_index"); + conf.set(InferenceConfig.MODEL_NAME, InMemorySearchFixture.MODEL_NAME); + return conf; + } + + private static TableDocument tableDoc(IndexMapping mapping, String id, String comment) { + String[] parts = id.split("\\."); + return MetastoreTableMapper.fromTable( + InMemorySearchFixture.table(parts[0], parts[1], parts[2], comment), mapping); + } + + private static final class IndexerFixture implements AutoCloseable { + private final IndexManager indexManager; + private final Indexer indexer; + private final EmbedModelRegistry modelRegistry; + + private IndexerFixture(IndexManager indexManager, Indexer indexer, EmbedModelRegistry modelRegistry) { + this.indexManager = indexManager; + this.indexer = indexer; + this.modelRegistry = modelRegistry; + } + + static IndexerFixture create(Configuration conf) throws IOException { + IndexMapping mapping = MetastoreSchemas.defaultHiveTablesMapping( + "test_index", InMemorySearchFixture.MODEL_NAME, conf); + IndexManager indexManager = IndexManager.open(mapping, conf); + EmbedModelRegistry registry = new EmbedModelRegistry( + Map.of(InMemorySearchFixture.MODEL_NAME, + new StubEmbedModel(InMemorySearchFixture.MODEL_NAME))); + Indexer indexer = new Indexer(indexManager, registry); + indexer.initialize(); + return new IndexerFixture(indexManager, indexer, registry); + } + + IndexMapping mapping() { + return indexManager.mapping(); + } + + @Override + public void close() throws Exception { + indexer.close(); + indexManager.close(); + modelRegistry.close(); + } + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/index/TestIndexerSearchIntegration.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/index/TestIndexerSearchIntegration.java new file mode 100644 index 000000000000..ffbc33c96061 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/index/TestIndexerSearchIntegration.java @@ -0,0 +1,123 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.index; + +import org.apache.hadoop.hive.common.TableName; +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hadoop.hive.metastore.api.Table; +import org.apache.hive.search.search.InMemorySearchFixture; +import org.apache.hive.search.search.TableSearchHit; +import org.apache.hive.search.search.TableSearchResult; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.util.List; +import java.util.Map; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +@Category(MetastoreUnitTest.class) +public class TestIndexerSearchIntegration { + + @Test + public void keywordSearchFindsIndexedTable() throws Exception { + try (InMemorySearchFixture fixture = InMemorySearchFixture.create()) { + fixture.mutations().addTable( + InMemorySearchFixture.table("hive", "sales", "orders", "daily sales orders")); + fixture.mutations().addTable( + InMemorySearchFixture.table("hive", "inventory", "parts", "spare parts catalog")); + fixture.commit(1L); + + List hits = fixture.searchMatch("sales", 5); + assertFalse(hits.isEmpty()); + assertTrue(hits.stream().anyMatch(hit -> hit.table().toString().contains("sales.orders"))); + } + } + + @Test + public void tableKeywordMatchesTableNameWhenCommentDoesNot() throws Exception { + try (InMemorySearchFixture fixture = InMemorySearchFixture.create()) { + fixture.mutations().addTable( + InMemorySearchFixture.table("hive", "sales", "customers", "customer master data")); + fixture.commit(1L); + + assertFalse(fixture.searchMatch("customers", 5).isEmpty()); + } + } + + @Test + public void tableKeywordRanksTableNameAboveColumnNameAboveComment() throws Exception { + try (InMemorySearchFixture fixture = InMemorySearchFixture.create()) { + fixture.mutations().addTable( + InMemorySearchFixture.table("hive", "sales", "orders", "revenue summary")); + fixture.mutations().addTable(tableWithColumn("hive", "sales", "line_items", "revenue")); + fixture.mutations().addTable( + InMemorySearchFixture.table("hive", "sales", "revenue", "monthly facts")); + fixture.commit(1L); + + List hits = fixture.searchMatch("revenue", 5); + assertEquals(3, hits.size()); + assertTrue(hits.get(0).table().toString().contains("sales.revenue")); + assertTrue(hits.get(1).table().toString().contains("sales.line_items")); + assertTrue(hits.get(2).table().toString().contains("sales.orders")); + } + } + + private static Table tableWithColumn(String catalog, String db, String name, String columnName) { + Table table = InMemorySearchFixture.table(catalog, db, name, "unrelated comment"); + table.getSd().setCols(List.of( + new org.apache.hadoop.hive.metastore.api.FieldSchema(columnName, "double", "amount value"))); + return table; + } + + @Test + public void incrementalDropRemovesTableFromSearchResults() throws Exception { + try (InMemorySearchFixture fixture = InMemorySearchFixture.create()) { + fixture.mutations().addTable( + InMemorySearchFixture.table("hive", "sales", "orders", "daily sales orders")); + fixture.commit(1L); + assertFalse(fixture.searchMatch("sales", 5).isEmpty()); + + fixture.mutations().dropTable(new TableName("hive", "sales", "orders")); + fixture.commit(2L); + + List hits = fixture.searchMatch("sales", 5); + assertTrue(hits.isEmpty()); + } + } + + @Test + public void alterTableUpdatesSearchText() throws Exception { + try (InMemorySearchFixture fixture = InMemorySearchFixture.create()) { + var before = InMemorySearchFixture.table("hive", "sales", "orders", "old comment"); + var after = InMemorySearchFixture.table("hive", "sales", "orders", "revenue analytics"); + fixture.mutations().addTable(before); + fixture.commit(1L); + assertTrue(fixture.searchMatch("revenue", 5).isEmpty()); + + fixture.mutations().replaceTable(before, after); + fixture.commit(2L); + + List hits = fixture.searchMatch("revenue", 5); + assertEquals(1, hits.size()); + assertTrue(hits.get(0).table().toString().contains("sales.orders")); + } + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/index/manifest/TestIndexManifest.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/index/manifest/TestIndexManifest.java new file mode 100644 index 000000000000..3edb4e851252 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/index/manifest/TestIndexManifest.java @@ -0,0 +1,122 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.index.manifest; + +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +@Category(MetastoreUnitTest.class) +public class TestIndexManifest { + + @Test + public void jsonRoundTrip() throws Exception { + IndexManifest manifest = IndexManifest.create( + "hive_tables", + List.of( + new IndexManifest.IndexFile("segments_1", 100), + new IndexManifest.IndexFile("segments_2", 200)), + "bge-small", + 42L); + IndexManifest parsed = IndexManifest.fromJson(manifest.toJsonBytes()); + assertEquals(manifest, parsed); + } + + @Test + public void sameFilesAsIgnoresOrder() { + IndexManifest left = IndexManifest.create( + "idx", + List.of( + new IndexManifest.IndexFile("a", 1), + new IndexManifest.IndexFile("b", 2)), + "model", + 1L); + IndexManifest right = IndexManifest.create( + "idx", + List.of( + new IndexManifest.IndexFile("b", 2), + new IndexManifest.IndexFile("a", 1)), + "model", + 2L); + assertTrue(left.sameFilesAs(right)); + } + + @Test + public void diffDetectsAddsUpdatesAndDeletes() { + IndexManifest source = IndexManifest.create( + "idx", + List.of( + new IndexManifest.IndexFile("keep", 10), + new IndexManifest.IndexFile("changed", 20), + new IndexManifest.IndexFile("added", 30)), + "model", + 5L); + IndexManifest target = IndexManifest.create( + "idx", + List.of( + new IndexManifest.IndexFile("keep", 10), + new IndexManifest.IndexFile("changed", 99), + new IndexManifest.IndexFile("removed", 40)), + "model", + 4L); + + List ops = source.diff(target); + assertEquals(3, ops.size()); + assertTrue(ops.contains(IndexManifest.ChangedFileOp.add("changed", 20L))); + assertTrue(ops.contains(IndexManifest.ChangedFileOp.add("added", 30L))); + assertTrue(ops.contains(IndexManifest.ChangedFileOp.del("removed"))); + } + + @Test + public void diffFromEmptyTargetAddsAllFiles() { + IndexManifest source = IndexManifest.create( + "idx", + List.of(new IndexManifest.IndexFile("segments_1", 100)), + "model", + 1L); + + List ops = source.diff(null); + assertEquals(1, ops.size()); + assertEquals(new IndexManifest.ChangedFileOp.Add("segments_1", 100L), ops.get(0)); + } + + @Test + public void diffAgainstIdenticalManifestIsEmpty() { + IndexManifest manifest = IndexManifest.create( + "idx", + List.of(new IndexManifest.IndexFile("segments_1", 100)), + "model", + 1L); + assertTrue(manifest.diff(manifest).isEmpty()); + } + + @Test + public void sameFilesAsReturnsFalseForDifferentSizes() { + IndexManifest left = IndexManifest.create( + "idx", List.of(new IndexManifest.IndexFile("a", 1)), "model", 1L); + IndexManifest right = IndexManifest.create( + "idx", List.of(new IndexManifest.IndexFile("a", 2)), "model", 1L); + assertFalse(left.sameFilesAs(right)); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/index/store/TestIndexBackupFailureRecovery.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/index/store/TestIndexBackupFailureRecovery.java new file mode 100644 index 000000000000..ec5ae1b359f7 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/index/store/TestIndexBackupFailureRecovery.java @@ -0,0 +1,120 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.index.store; + +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hive.search.index.manifest.IndexManifest; +import org.apache.hive.search.testutil.InMemoryIndexStateClient; +import org.junit.Before; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.io.ByteArrayInputStream; +import java.util.List; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +@Category(MetastoreUnitTest.class) +public class TestIndexBackupFailureRecovery { + + private InMemoryIndexStateClient local; + private InMemoryIndexStateClient remote; + + @Before + public void setUp() { + local = new InMemoryIndexStateClient(); + remote = new InMemoryIndexStateClient(); + } + + @Test + public void restoreResumesAfterPartialFileCopy() throws Exception { + IndexManifest target = manifest(20L, "segments_1", 100, "segments_2", 200); + writeManifest(remote, target); + remote.write("segments_1", new ByteArrayInputStream(new byte[] {1})); + remote.write("segments_2", new ByteArrayInputStream(new byte[] {2})); + + local.write("segments_1", new ByteArrayInputStream(new byte[] {1})); + local.writeStagingManifest(target); + + assertTrue(IndexBackupUtils.restoreFromBackup(local, remote)); + assertTrue(local.hasFile("segments_2")); + assertFalse(local.readStagingManifest().isPresent()); + } + + @Test + public void resolveInterruptedRestoreKeepsStagingWhenFilesMismatch() throws Exception { + IndexManifest target = manifest(15L, "segments_1", 100, "segments_2", 200); + writeManifest(local, manifest(15L, "segments_1", 100)); + local.writeStagingManifest(target); + + IndexBackupUtils.resolveInterruptedRestore(local); + + assertTrue(local.readStagingManifest().isPresent()); + } + + @Test + public void syncToBackupRemovesStaleRemoteFiles() throws Exception { + writeManifest(remote, manifest(5L, "segments_old", 50)); + remote.write("segments_old", new ByteArrayInputStream(new byte[] {9})); + + writeManifest(local, manifest(10L, "segments_new", 100)); + local.write("segments_new", new ByteArrayInputStream(new byte[] {1, 2, 3})); + + assertTrue(IndexBackupUtils.syncToBackup(local, remote)); + assertFalse(remote.hasFile("segments_old")); + assertTrue(remote.hasFile("segments_new")); + } + + @Test + public void restoreFromBackupRemovesStaleLocalFiles() throws Exception { + writeManifest(local, manifest(5L, "segments_old", 50)); + local.write("segments_old", new ByteArrayInputStream(new byte[] {9})); + + writeManifest(remote, manifest(20L, "segments_new", 100)); + remote.write("segments_new", new ByteArrayInputStream(new byte[] {3, 2, 1})); + + assertTrue(IndexBackupUtils.restoreFromBackup(local, remote)); + assertFalse(local.hasFile("segments_old")); + assertTrue(local.hasFile("segments_new")); + } + + private static IndexManifest manifest(long eventId, String fileName, long size) { + return IndexManifest.create( + "hive_tables", + List.of(new IndexManifest.IndexFile(fileName, size)), + "bge-small", + eventId); + } + + private static IndexManifest manifest( + long eventId, String fileName1, long size1, String fileName2, long size2) { + return IndexManifest.create( + "hive_tables", + List.of( + new IndexManifest.IndexFile(fileName1, size1), + new IndexManifest.IndexFile(fileName2, size2)), + "bge-small", + eventId); + } + + private static void writeManifest(InMemoryIndexStateClient client, IndexManifest manifest) + throws Exception { + client.write(IndexManifest.MANIFEST_FILE_NAME, new ByteArrayInputStream(manifest.toJsonBytes())); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/index/store/TestIndexBackupUtils.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/index/store/TestIndexBackupUtils.java new file mode 100644 index 000000000000..8fcf73c5ea4f --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/index/store/TestIndexBackupUtils.java @@ -0,0 +1,111 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.index.store; + +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hive.search.index.manifest.IndexManifest; +import org.apache.hive.search.testutil.InMemoryIndexStateClient; +import org.junit.Before; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.io.ByteArrayInputStream; +import java.util.List; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +@Category(MetastoreUnitTest.class) +public class TestIndexBackupUtils { + + private InMemoryIndexStateClient local; + private InMemoryIndexStateClient remote; + + @Before + public void setUp() { + local = new InMemoryIndexStateClient(); + remote = new InMemoryIndexStateClient(); + } + + @Test + public void syncToBackupCopiesNewManifestAndFiles() throws Exception { + writeManifest(local, manifest(10L, "segments_1", 100)); + local.write("segments_1", new ByteArrayInputStream(new byte[] {1, 2, 3})); + + assertTrue(IndexBackupUtils.syncToBackup(local, remote)); + assertTrue(remote.hasFile(IndexManifest.MANIFEST_FILE_NAME)); + assertTrue(remote.hasFile("segments_1")); + assertEqualsEventId(10L, remote); + } + + @Test + public void syncToBackupSkipsWhenRemoteIsNewer() throws Exception { + writeManifest(local, manifest(5L, "segments_1", 100)); + writeManifest(remote, manifest(10L, "segments_1", 100)); + + assertFalse(IndexBackupUtils.syncToBackup(local, remote)); + } + + @Test + public void restoreFromBackupPullsRemoteState() throws Exception { + writeManifest(remote, manifest(20L, "segments_2", 200)); + remote.write("segments_2", new ByteArrayInputStream(new byte[] {9, 8, 7})); + + assertTrue(IndexBackupUtils.restoreFromBackup(local, remote)); + assertTrue(local.hasFile("segments_2")); + assertFalse(local.readStagingManifest().isPresent()); + } + + @Test + public void restoreFromBackupSkipsWhenLocalIsNewer() throws Exception { + writeManifest(local, manifest(30L, "segments_1", 100)); + writeManifest(remote, manifest(10L, "segments_1", 100)); + + assertFalse(IndexBackupUtils.restoreFromBackup(local, remote)); + } + + @Test + public void resolveInterruptedRestoreClearsMatchingStagingManifest() throws Exception { + IndexManifest target = manifest(15L, "segments_3", 300); + writeManifest(local, target); + local.writeStagingManifest(target); + + IndexBackupUtils.resolveInterruptedRestore(local); + + assertFalse(local.readStagingManifest().isPresent()); + } + + private static IndexManifest manifest(long eventId, String fileName, long size) { + return IndexManifest.create( + "hive_tables", + List.of(new IndexManifest.IndexFile(fileName, size)), + "bge-small", + eventId); + } + + private static void writeManifest(InMemoryIndexStateClient client, IndexManifest manifest) + throws Exception { + client.write(IndexManifest.MANIFEST_FILE_NAME, new ByteArrayInputStream(manifest.toJsonBytes())); + } + + private static void assertEqualsEventId(long expected, InMemoryIndexStateClient client) + throws Exception { + assertTrue(client.readManifest().isPresent()); + assertTrue(expected == client.readManifest().get().lastEventId()); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/inference/TestEmbedModelRegistry.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/inference/TestEmbedModelRegistry.java new file mode 100644 index 000000000000..fb95f4836093 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/inference/TestEmbedModelRegistry.java @@ -0,0 +1,61 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.inference; + +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hive.search.exception.IndexException; +import org.apache.hive.search.testutil.StubEmbedModel; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.util.Map; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +@Category(MetastoreUnitTest.class) +public class TestEmbedModelRegistry { + + @Test + public void getReturnsConfiguredModel() throws Exception { + StubEmbedModel model = new StubEmbedModel("stub-model"); + EmbedModelRegistry registry = new EmbedModelRegistry(Map.of("stub-model", model)); + assertTrue(registry.get("stub-model") == model); + registry.close(); + } + + @Test + public void getThrowsForUnknownModel() { + EmbedModelRegistry registry = new EmbedModelRegistry(Map.of()); + IndexException error = assertThrows(IndexException.class, () -> registry.get("missing")); + assertTrue(error.getMessage().contains("missing")); + } + + @Test + public void stubModelProducesDeterministicVectors() throws Exception { + StubEmbedModel model = new StubEmbedModel("stub-model"); + float[] first = model.embed(EmbedModel.TaskType.QUERY, "sales"); + float[] second = model.embed(EmbedModel.TaskType.QUERY, "sales"); + assertArrayEquals(first, second, 0.0001f); + + float[] document = model.embed(EmbedModel.TaskType.DOCUMENT, "sales"); + assertNotEquals(first[0], document[0], 0.0001f); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/inference/TestEmbeddingCache.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/inference/TestEmbeddingCache.java new file mode 100644 index 000000000000..9ef2ca0ffbf2 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/inference/TestEmbeddingCache.java @@ -0,0 +1,87 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.inference; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hive.search.config.InferenceConfig; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +@Category(MetastoreUnitTest.class) +public class TestEmbeddingCache { + + @Test + public void getReturnsCopyOfStoredVector() { + EmbeddingCache cache = EmbeddingCache.create(enabledCacheConf()); + float[] embedding = {0.1f, 0.2f, 0.3f}; + cache.put("model-a", EmbedModel.TaskType.DOCUMENT, "sales orders", embedding); + + float[] cached = cache.get("model-a", EmbedModel.TaskType.DOCUMENT, "sales orders").orElseThrow(); + assertArrayEquals(embedding, cached, 0.001f); + cached[0] = 9f; + float[] again = cache.get("model-a", EmbedModel.TaskType.DOCUMENT, "sales orders").orElseThrow(); + assertEquals(0.1f, again[0], 0.001f); + } + + @Test + public void cacheIsScopedByModelAndTask() { + EmbeddingCache cache = EmbeddingCache.create(enabledCacheConf()); + float[] doc = {1f}; + float[] query = {2f}; + cache.put("model-a", EmbedModel.TaskType.DOCUMENT, "text", doc); + cache.put("model-a", EmbedModel.TaskType.QUERY, "text", query); + cache.put("model-b", EmbedModel.TaskType.DOCUMENT, "text", new float[] {3f}); + + assertArrayEquals(doc, cache.get("model-a", EmbedModel.TaskType.DOCUMENT, "text").orElseThrow(), 0.001f); + assertArrayEquals(query, cache.get("model-a", EmbedModel.TaskType.QUERY, "text").orElseThrow(), 0.001f); + assertArrayEquals(new float[] {3f}, + cache.get("model-b", EmbedModel.TaskType.DOCUMENT, "text").orElseThrow(), 0.001f); + } + + @Test + public void disabledCacheAlwaysMisses() { + EmbeddingCache cache = EmbeddingCache.disabled(); + cache.put("model-a", EmbedModel.TaskType.DOCUMENT, "text", new float[] {1f}); + assertFalse(cache.get("model-a", EmbedModel.TaskType.DOCUMENT, "text").isPresent()); + assertEquals(1, cache.misses()); + } + + @Test + public void tracksHitsAndMisses() { + EmbeddingCache cache = EmbeddingCache.create(enabledCacheConf()); + cache.get("model-a", EmbedModel.TaskType.DOCUMENT, "missing"); + + cache.put("model-a", EmbedModel.TaskType.DOCUMENT, "sales", new float[] {1f}); + assertTrue(cache.get("model-a", EmbedModel.TaskType.DOCUMENT, "sales").isPresent()); + assertEquals(1, cache.hits()); + assertEquals(1, cache.misses()); + } + + private static Configuration enabledCacheConf() { + Configuration conf = new Configuration(false); + conf.setBoolean(InferenceConfig.EMBEDDING_CACHE_ENABLED, true); + conf.setInt(InferenceConfig.EMBEDDING_CACHE_MAX_ENTRIES, 1000); + return conf; + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/inference/TestEmbeddingPrompt.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/inference/TestEmbeddingPrompt.java new file mode 100644 index 000000000000..3c1642e1962b --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/inference/TestEmbeddingPrompt.java @@ -0,0 +1,42 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.inference; + +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import static org.junit.Assert.assertEquals; + +@Category(MetastoreUnitTest.class) +public class TestEmbeddingPrompt { + + @Test + public void prefixForTaskType() { + EmbeddingPrompt prompt = EmbeddingPrompt.e5(); + assertEquals("passage: ", prompt.prefixFor(EmbedModel.TaskType.DOCUMENT)); + assertEquals("query: ", prompt.prefixFor(EmbedModel.TaskType.QUERY)); + } + + @Test + public void noneUsesEmptyPrefixes() { + EmbeddingPrompt prompt = EmbeddingPrompt.none(); + assertEquals("", prompt.prefixFor(EmbedModel.TaskType.DOCUMENT)); + assertEquals("", prompt.prefixFor(EmbedModel.TaskType.QUERY)); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/inference/TestLocalOnnxEmbeddingModel.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/inference/TestLocalOnnxEmbeddingModel.java new file mode 100644 index 000000000000..a027767bcc8b --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/inference/TestLocalOnnxEmbeddingModel.java @@ -0,0 +1,89 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.inference; + +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; + +@Category(MetastoreUnitTest.class) +public class TestLocalOnnxEmbeddingModel { + + @Test + public void meanPoolAveragesMaskedTokens() { + float[][] tokenEmbeddings = { + {1f, 0f}, + {3f, 0f}, + {9f, 9f} + }; + long[] mask = {1, 1, 0}; + + assertArrayEquals(new float[] {2f, 0f}, InferenceWorker.meanPool(tokenEmbeddings, mask, 3), 0.001f); + } + + @Test + public void normalizeProducesUnitVector() { + float[] vector = {3f, 4f}; + InferenceWorker.normalize(vector); + float norm = 0f; + for (float v : vector) { + norm += v * v; + } + assertEquals(1f, norm, 0.001f); + assertEquals(0.6f, vector[0], 0.001f); + assertEquals(0.8f, vector[1], 0.001f); + } + + @Test + public void meanPoolVectorsAveragesChunkEmbeddings() { + float[][] chunks = { + {1f, 0f}, + {3f, 0f} + }; + + assertArrayEquals(new float[] {2f, 0f}, InferenceWorker.meanPoolVectors(chunks), 0.001f); + } + + @Test + public void meanPoolVectorsReturnsSingleChunkUnchanged() { + float[][] chunks = {{0.6f, 0.8f}}; + + assertArrayEquals(new float[] {0.6f, 0.8f}, InferenceWorker.meanPoolVectors(chunks), 0.001f); + } + + @Test + public void chunkedMeanPoolThenNormalizeProducesUnitVector() { + float[][] chunks = { + {1f, 0f}, + {0f, 1f} + }; + float[] pooled = InferenceWorker.meanPoolVectors(chunks); + InferenceWorker.normalize(pooled); + + float norm = 0f; + for (float v : pooled) { + norm += v * v; + } + assertEquals(1f, norm, 0.001f); + assertEquals(0.7071f, pooled[0], 0.001f); + assertEquals(0.7071f, pooled[1], 0.001f); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/mapping/TestIndexMapping.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/mapping/TestIndexMapping.java new file mode 100644 index 000000000000..634c00839b71 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/mapping/TestIndexMapping.java @@ -0,0 +1,62 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.mapping; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hive.search.metastore.MetastoreSchemas; +import org.apache.hive.search.metastore.MetastoreTableMapper; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +@Category(MetastoreUnitTest.class) +public class TestIndexMapping { + + @Test + public void defaultSchemaExposesSingleHybridField() { + Configuration conf = new Configuration(false); + IndexMapping mapping = MetastoreSchemas.defaultHiveTablesMapping("hive_tables", "bge-small", conf); + + assertEquals(1, mapping.hybridFields().size()); + assertTrue(mapping.soleHybridField().isPresent()); + assertEquals(MetastoreTableMapper.FIELD_SEARCH_TEXT, mapping.soleHybridField().get()); + } + + @Test + public void lexicalFieldsAreConfiguredInDefaultSchema() { + Configuration conf = new Configuration(false); + IndexMapping mapping = MetastoreSchemas.defaultHiveTablesMapping("hive_tables", "bge-small", conf); + FieldSchema tableSchema = mapping.fieldSchema(MetastoreTableMapper.FIELD_TABLE); + FieldSchema commentSchema = mapping.fieldSchema(MetastoreTableMapper.FIELD_COMMENT); + assertTrue(tableSchema instanceof FieldSchema.TextFieldSchema table + && table.search().lexical()); + assertTrue(commentSchema instanceof FieldSchema.TextFieldSchema comment + && comment.search().lexical()); + FieldSchema columnNamesSchema = mapping.fieldSchema(MetastoreTableMapper.FIELD_COLUMN_NAMES); + FieldSchema columnCommentsSchema = mapping.fieldSchema(MetastoreTableMapper.FIELD_COLUMN_COMMENTS); + assertTrue(columnNamesSchema instanceof FieldSchema.TextFieldSchema columnNames + && columnNames.search().lexical()); + assertTrue(columnCommentsSchema instanceof FieldSchema.TextFieldSchema columnComments + && columnComments.search().lexical()); + assertNotNull(mapping.analyzer()); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/RealMetastoreSearchSession.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/RealMetastoreSearchSession.java new file mode 100644 index 000000000000..e4ede26f6f5d --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/RealMetastoreSearchSession.java @@ -0,0 +1,250 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.metastore; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.metastore.HiveMetaStoreClient; +import org.apache.hadoop.hive.metastore.MetaStoreTestUtils; +import org.apache.hadoop.hive.metastore.conf.MetastoreConf; +import org.apache.hadoop.hive.metastore.conf.MetastoreConf.ConfVars; +import org.apache.hadoop.hive.metastore.messaging.json.JSONMessageEncoder; +import org.apache.hive.search.config.IndexConfig; +import org.apache.hive.search.config.IndexStoreConfig; +import org.apache.hive.search.config.InferenceConfig; +import org.apache.hive.search.config.SearchConfig; +import org.apache.hive.search.index.Indexer; +import org.apache.hive.search.index.IndexManager; +import org.apache.hive.search.index.store.LocalStateClient; +import org.apache.hive.search.inference.EmbedModelRegistry; +import org.apache.hive.search.mapping.IndexMapping; +import org.apache.hive.search.search.SearchArgs; +import org.apache.hive.search.search.SearchInternal; +import org.apache.hive.search.search.SearchQuery; +import org.apache.hive.search.search.TableSearchResult; +import org.apache.hive.search.testutil.InMemoryIndexStateClient; +import org.apache.hive.search.testutil.RealMetastoreServer; +import org.apache.hive.search.testutil.StubEmbedModel; +import org.apache.hive.search.testutil.TestLeaderElection; +import org.apache.lucene.search.BayesianScoreEstimator; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.SearcherManager; +import org.apache.lucene.store.ByteBuffersDirectory; + +import java.io.IOException; +import java.util.List; +import java.util.Map; +import java.util.UUID; + +/** End-to-end search session wired to a real HMS and in-memory Lucene index. */ +public final class RealMetastoreSearchSession implements AutoCloseable { + public static final String MODEL_NAME = "stub-model"; + private static final String ELECTION = "real-metastore-search-test"; + + private final IndexManager indexManager; + private final Indexer indexer; + private final EmbedModelRegistry modelRegistry; + private final MetastoreIndexer metastoreIndexer; + private final SearchConfig searchConfig; + + private SearcherManager searcherManager; + private BayesianScoreEstimator.Parameters bayesianParameters; + + private RealMetastoreSearchSession( + IndexManager indexManager, + Indexer indexer, + EmbedModelRegistry modelRegistry, + MetastoreIndexer metastoreIndexer, + SearchConfig searchConfig) { + this.indexManager = indexManager; + this.indexer = indexer; + this.modelRegistry = modelRegistry; + this.metastoreIndexer = metastoreIndexer; + this.searchConfig = searchConfig; + } + + public static RealMetastoreSearchSession open(RealMetastoreServer server) throws Exception { + return open(server, null); + } + + public static RealMetastoreSearchSession open( + RealMetastoreServer server, InMemoryIndexStateClient sharedRemote) throws Exception { + Configuration conf = searchConfiguration(server.conf()); + IndexMapping mapping = MetastoreSchemas.defaultHiveTablesMapping( + "test_index", MODEL_NAME, conf); + + ByteBuffersDirectory directory = new ByteBuffersDirectory(); + LocalStateClient local = new LocalStateClient(directory, "test_index"); + IndexManager indexManager = sharedRemote == null + ? new IndexManager(mapping, directory, local, null) + : new IndexManager(mapping, directory, local, sharedRemote); + + EmbedModelRegistry modelRegistry = new EmbedModelRegistry( + Map.of(MODEL_NAME, new StubEmbedModel(MODEL_NAME))); + Indexer indexer = new Indexer(indexManager, modelRegistry); + HiveMetaStoreClient sessionClient = new HiveMetaStoreClient(server.conf()); + MetastoreIndexer metastoreIndexer = + new MetastoreIndexer(conf, indexManager, indexer, sessionClient, false); + return new RealMetastoreSearchSession( + indexManager, indexer, modelRegistry, metastoreIndexer, new SearchConfig(conf)); + } + + public static InMemoryIndexStateClient newSharedRemoteBackup() { + return new InMemoryIndexStateClient(); + } + + public static String uniqueDbName(String prefix) { + return prefix + "_" + UUID.randomUUID().toString().replace('-', '_').substring(0, 8); + } + + private static Configuration searchConfiguration(Configuration hmsConf) { + Configuration conf = new Configuration(false); + conf.setBoolean(IndexStoreConfig.MEMORY, true); + conf.set(IndexConfig.INDEX_NAME, "test_index"); + conf.set(InferenceConfig.MODEL_NAME, MODEL_NAME); + conf.setInt(IndexConfig.BOOTSTRAP_FETCH_THREADS, 2); + conf.setInt(IndexConfig.BOOTSTRAP_QUEUE_DEPTH, 8); + conf.setLong(IndexConfig.BOOTSTRAP_PROGRESS_INTERVAL_MS, Long.MAX_VALUE); + conf.setLong(IndexConfig.EVENT_FAILURE_BACKOFF_MS, 0L); + conf.setLong(IndexConfig.FLUSH_INTERVAL_MS, 60_000L); + conf.setLong(IndexConfig.INDEX_SYNC_INTERVAL, 60_000L); + conf.setInt(SearchConfig.BAYESIAN_SAMPLES, 8); + conf.setInt(SearchConfig.BAYESIAN_TOKENS_PER_QUERY, 3); + conf.setLong(SearchConfig.BAYESIAN_SEED, 1L); + MetastoreConf.setVar(conf, ConfVars.EVENT_MESSAGE_FACTORY, JSONMessageEncoder.class.getName()); + MetastoreConf.setBoolVar(conf, ConfVars.CAPABILITY_CHECK, false); + MetastoreConf.setVar(conf, ConfVars.THRIFT_URIS, MetastoreConf.getVar(hmsConf, ConfVars.THRIFT_URIS)); + MetastoreConf.setVar(conf, ConfVars.THRIFT_BIND_HOST, MetastoreConf.getVar(hmsConf, ConfVars.THRIFT_BIND_HOST)); + TestLeaderElection election = new TestLeaderElection(true); + MetastoreCluster.injectElection(conf, ELECTION, election, "mutex"); + return conf; + } + + public void drainNotifications() throws Exception { + for (int round = 0; round < 20; round++) { + if (metastoreIndexer.pollEvents(100) == 0) { + break; + } + Thread.sleep(50); + } + metastoreIndexer.flushCheckpoint(); + refreshSearcher(); + } + + public void waitUntilSearchable(String text, int limit) throws Exception { + MetaStoreTestUtils.waitForAssertion( + "table indexed for query '" + text + "'", + () -> { + try { + drainNotifications(); + if (searchMatch(text, limit).hits().isEmpty()) { + throw new AssertionError("not searchable yet"); + } + } catch (AssertionError e) { + throw e; + } catch (Exception e) { + throw new RuntimeException(e); + } + }, + 30_000, + 300); + } + + public void waitUntilNotSearchable(String text, int limit) throws Exception { + MetaStoreTestUtils.waitForAssertion( + "table removed for query '" + text + "'", + () -> { + try { + drainNotifications(); + if (!searchMatch(text, limit).hits().isEmpty()) { + throw new AssertionError("still searchable for '" + text + "'"); + } + } catch (RuntimeException e) { + throw e; + } catch (Exception e) { + throw new RuntimeException(e); + } + }, + 30_000, + 300); + } + + public void syncBackup() throws IOException { + metastoreIndexer.syncBackup(); + } + + public void refreshSearcher() throws IOException { + if (searcherManager == null) { + searcherManager = new SearcherManager(indexer.writer(), null); + } else { + searcherManager.maybeRefresh(); + } + if (bayesianParameters == null) { + IndexSearcher searcher = searcherManager.acquire(); + try { + bayesianParameters = BayesianScoreEstimator.estimate( + searcher, + MetastoreTableMapper.FIELD_SEARCH_TEXT, + searchConfig.getBayesianSamples(), + searchConfig.getBayesianTokensPerQuery(), + searchConfig.getBayesianSeed()); + } finally { + searcherManager.release(searcher); + } + } + } + + public TableSearchResult searchMatch(String text, int limit) throws Exception { + refreshSearcher(); + try (SearchInternal searchIO = new SearchInternal( + searcherManager, indexManager, modelRegistry, searchConfig, bayesianParameters)) { + return searchIO.search(new SearchQuery( + new SearchArgs.Match(text), + null, null, limit, + List.of(MetastoreTableMapper.FIELD_TABLE, MetastoreTableMapper.FIELD_COMMENT))); + } + } + + public TableSearchResult searchHybrid(String queryText, int limit) throws Exception { + refreshSearcher(); + try (SearchInternal searchIO = new SearchInternal( + searcherManager, indexManager, modelRegistry, searchConfig, bayesianParameters)) { + return searchIO.search(SearchQuery.fromQueryBody( + Map.of( + "field", MetastoreTableMapper.FIELD_SEARCH_TEXT, + "query", queryText), + SearchQuery.Mode.HYBRID, + null, null, limit, + List.of(MetastoreTableMapper.FIELD_TABLE, MetastoreTableMapper.FIELD_COMMENT))); + } + } + + public IndexManager indexManager() { + return indexManager; + } + + @Override + public void close() throws Exception { + if (searcherManager != null) { + searcherManager.close(); + } + metastoreIndexer.close(); + indexer.close(); + indexManager.close(); + modelRegistry.close(); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestBootstrapIndexer.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestBootstrapIndexer.java new file mode 100644 index 000000000000..5852660b1ec2 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestBootstrapIndexer.java @@ -0,0 +1,172 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.metastore; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hadoop.hive.metastore.api.Table; +import org.apache.hive.search.config.IndexConfig; +import org.apache.hive.search.config.IndexStoreConfig; +import org.apache.hive.search.exception.IndexException; +import org.apache.hive.search.index.Indexer; +import org.apache.hive.search.index.IndexManager; +import org.apache.hive.search.inference.EmbedModelRegistry; +import org.apache.hadoop.hive.metastore.IMetaStoreClient; +import org.apache.hive.search.mapping.IndexMapping; +import org.apache.hive.search.search.InMemorySearchFixture; +import org.apache.hive.search.testutil.MetastoreBootstrapMocks; +import org.apache.hive.search.testutil.StubEmbedModel; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.util.List; +import java.util.Map; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.anyList; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@Category(MetastoreUnitTest.class) +public class TestBootstrapIndexer { + + @Test + public void bootstrapsAllTablesAcrossDatabases() throws Exception { + Configuration conf = bootstrapConf(); + Table orders = InMemorySearchFixture.table("hive", "sales", "orders", "sales orders"); + Table customers = InMemorySearchFixture.table("hive", "sales", "customers", "sales customers"); + Table parts = InMemorySearchFixture.table("hive", "inventory", "parts", "spare parts"); + + try (BootstrapFixture fixture = BootstrapFixture.create(conf)) { + MetastoreBootstrapMocks.stubBootstrapCatalog(fixture.client(), orders, customers, parts); + + new BootstrapIndexer(conf, fixture.mapping(), fixture.indexer(), fixture.client(), true) + .run(500L); + + assertEquals(500L, fixture.indexManager().readLocalManifest().orElseThrow().lastEventId()); + assertEquals(3, fixture.indexer().writer().getDocStats().numDocs); + } + } + + @Test + public void bootstrapHonorsBatchSize() throws Exception { + Configuration conf = bootstrapConf(); + conf.setInt(IndexConfig.BOOTSTRAP_BATCH_SIZE, 1); + conf.setInt(IndexConfig.BOOTSTRAP_FETCH_THREADS, 1); + Table orders = InMemorySearchFixture.table("hive", "sales", "orders", "sales orders"); + Table customers = InMemorySearchFixture.table("hive", "sales", "customers", "sales customers"); + + try (BootstrapFixture fixture = BootstrapFixture.create(conf)) { + MetastoreBootstrapMocks.stubBootstrapCatalog(fixture.client(), orders, customers); + + new BootstrapIndexer(conf, fixture.mapping(), fixture.indexer(), fixture.client(), true) + .run(10L); + + assertEquals(2, fixture.indexer().writer().getDocStats().numDocs); + } + } + + @Test + public void bootstrapFailsWhenMetastoreFetchFails() throws Exception { + Configuration conf = bootstrapConf(); + try (BootstrapFixture fixture = BootstrapFixture.create(conf)) { + IMetaStoreClient client = fixture.client(); + when(client.getAllDatabases()).thenReturn(List.of("sales")); + when(client.getAllTables("sales")).thenReturn(List.of("orders")); + when(client.getTableObjectsByName(eq("sales"), anyList())) + .thenThrow(new RuntimeException("metastore unavailable")); + + IndexException error = assertThrows(IndexException.class, + () -> new BootstrapIndexer(conf, fixture.mapping(), fixture.indexer(), client, true) + .run(1L)); + assertTrue(error.getMessage().contains("Bootstrap indexing failed") + || error.getMessage().contains("metastore unavailable")); + } + } + + private static Configuration bootstrapConf() { + Configuration conf = new Configuration(false); + conf.setInt(IndexConfig.BOOTSTRAP_FETCH_THREADS, 1); + conf.setInt(IndexConfig.BOOTSTRAP_QUEUE_DEPTH, 4); + conf.setInt(IndexConfig.COMMIT_FLUSHES, 1); + conf.setLong(IndexConfig.BOOTSTRAP_PROGRESS_INTERVAL_MS, Long.MAX_VALUE); + return conf; + } + + private static final class BootstrapFixture implements AutoCloseable { + private final IndexManager indexManager; + private final Indexer indexer; + private final EmbedModelRegistry modelRegistry; + private final IMetaStoreClient client; + + private BootstrapFixture( + IndexManager indexManager, + Indexer indexer, + EmbedModelRegistry modelRegistry, + IMetaStoreClient client) { + this.indexManager = indexManager; + this.indexer = indexer; + this.modelRegistry = modelRegistry; + this.client = client; + } + + static BootstrapFixture create(Configuration conf) throws Exception { + conf = new Configuration(conf); + conf.setBoolean(IndexStoreConfig.MEMORY, true); + conf.set(IndexConfig.INDEX_NAME, "test_index"); + conf.set(org.apache.hive.search.config.InferenceConfig.MODEL_NAME, + InMemorySearchFixture.MODEL_NAME); + org.apache.hive.search.mapping.IndexMapping mapping = + MetastoreSchemas.defaultHiveTablesMapping( + "test_index", InMemorySearchFixture.MODEL_NAME, conf); + IndexManager indexManager = IndexManager.open(mapping, conf); + EmbedModelRegistry registry = new EmbedModelRegistry( + Map.of(InMemorySearchFixture.MODEL_NAME, + new StubEmbedModel(InMemorySearchFixture.MODEL_NAME))); + Indexer indexer = new Indexer(indexManager, registry); + indexer.initialize(); + return new BootstrapFixture(indexManager, indexer, registry, mock(IMetaStoreClient.class)); + } + + IndexManager indexManager() { + return indexManager; + } + + Indexer indexer() { + return indexer; + } + + IndexMapping mapping() { + return indexManager.mapping(); + } + + IMetaStoreClient client() { + return client; + } + + @Override + public void close() throws Exception { + indexer.close(); + indexManager.close(); + modelRegistry.close(); + } + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestMetastoreEventHandlerFailureRecovery.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestMetastoreEventHandlerFailureRecovery.java new file mode 100644 index 000000000000..a9b2d0568163 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestMetastoreEventHandlerFailureRecovery.java @@ -0,0 +1,258 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.metastore; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.metastore.IMetaStoreClient; +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hadoop.hive.metastore.conf.MetastoreConf; +import org.apache.hadoop.hive.metastore.messaging.json.JSONMessageEncoder; +import org.apache.hadoop.hive.metastore.api.NotificationEvent; +import org.apache.hadoop.hive.metastore.api.NotificationEventRequest; +import org.apache.hadoop.hive.metastore.api.NotificationEventResponse; +import org.apache.hadoop.hive.metastore.api.Table; +import org.apache.hadoop.hive.metastore.messaging.MessageBuilder; +import org.apache.hive.search.config.IndexConfig; +import org.apache.hive.search.exception.IndexNotHealthyException; +import org.apache.hive.search.search.InMemorySearchFixture; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@Category(MetastoreUnitTest.class) +public class TestMetastoreEventHandlerFailureRecovery { + + @Test + public void batchFailureFallsBackToSingleEventApply() throws Exception { + Configuration conf = eventHandlerConf(); + conf.setInt(IndexConfig.EVENT_BATCH_MAX_FAILURES, 1); + + MessageBuilder messageBuilder = MessageBuilder.getInstance(); + Table orders = InMemorySearchFixture.table("hive", "sales", "orders", "sales orders"); + Table customers = InMemorySearchFixture.table("hive", "sales", "customers", "sales customers"); + NotificationEvent event1 = createEvent(101L, orders, messageBuilder); + NotificationEvent event2 = createEvent(102L, customers, messageBuilder); + + IMetaStoreClient client = mock(IMetaStoreClient.class); + when(client.getNextNotification(any(NotificationEventRequest.class), eq(true), eq(null))) + .thenReturn(response(event1, event2)) + .thenReturn(emptyResponse()); + + try (InMemorySearchFixture fixture = InMemorySearchFixture.create()) { + AtomicInteger batchAttempts = new AtomicInteger(); + MetastoreEventHandler handler = new MetastoreEventHandler(conf, client); + handler.addListeners(new MetastoreEventListener() { + @Override + public void notifyIndexTask(IndexTask task) throws IOException { + if (task.tablesToAdd.size() > 1) { + batchAttempts.incrementAndGet(); + throw new IOException("simulated batch apply failure"); + } + try { + fixture.mutations().apply(task); + fixture.commit(task.lastEventId); + } catch (Exception e) { + throw new IOException(e); + } + } + }); + + assertEquals(2, handler.getNextMetastoreEvents(10)); + assertEquals(1, batchAttempts.get()); + assertFalse(fixture.searchMatch("orders", 5).isEmpty()); + assertFalse(fixture.searchMatch("customers", 5).isEmpty()); + } + } + + @Test + public void poisonEventMarksIndexUnhealthyAfterPartialApply() throws Exception { + Configuration conf = eventHandlerConf(); + MessageBuilder messageBuilder = MessageBuilder.getInstance(); + Table orders = InMemorySearchFixture.table("hive", "sales", "orders", "sales orders"); + NotificationEvent good = createEvent(201L, orders, messageBuilder); + NotificationEvent poison = event( + 202L, + MessageBuilder.CREATE_TABLE_EVENT, + "{invalid-create-table-message", + "hive", + "sales", + "broken"); + + IMetaStoreClient client = mock(IMetaStoreClient.class); + when(client.getNextNotification(any(NotificationEventRequest.class), eq(true), eq(null))) + .thenReturn(response(good, poison)); + + try (InMemorySearchFixture fixture = InMemorySearchFixture.create()) { + MetastoreEventHandler handler = new MetastoreEventHandler(conf, client); + handler.addListeners(new MetastoreEventListener() { + @Override + public void notifyIndexTask(IndexTask task) throws IOException { + try { + fixture.mutations().apply(task); + fixture.commit(task.lastEventId); + } catch (Exception e) { + throw new IOException(e); + } + } + }); + + IndexNotHealthyException error = + assertThrows(IndexNotHealthyException.class, () -> handler.getNextMetastoreEvents(10)); + assertTrue(error.getMessage().contains("202")); + assertTrue(error.getMessage().contains("201")); + assertFalse(fixture.searchMatch("orders", 5).isEmpty()); + } + } + + @Test + public void skipPoisonEventAdvancesCursorAndRecovers() throws Exception { + Configuration conf = eventHandlerConf(); + conf.setBoolean(IndexConfig.EVENT_SKIP_POISON, true); + + MessageBuilder messageBuilder = MessageBuilder.getInstance(); + Table orders = InMemorySearchFixture.table("hive", "sales", "orders", "sales orders"); + NotificationEvent good = createEvent(301L, orders, messageBuilder); + NotificationEvent poison = event( + 302L, + MessageBuilder.CREATE_TABLE_EVENT, + "{invalid-create-table-message", + "hive", + "sales", + "broken"); + + IMetaStoreClient client = mock(IMetaStoreClient.class); + when(client.getNextNotification(any(NotificationEventRequest.class), eq(true), eq(null))) + .thenReturn(response(good, poison)) + .thenReturn(emptyResponse()); + + try (InMemorySearchFixture fixture = InMemorySearchFixture.create()) { + MetastoreEventHandler handler = new MetastoreEventHandler(conf, client); + handler.addListeners(new MetastoreEventListener() { + @Override + public void notifyIndexTask(IndexTask task) throws IOException { + try { + fixture.mutations().apply(task); + fixture.commit(task.lastEventId); + } catch (Exception e) { + throw new IOException(e); + } + } + }); + + assertEquals(2, handler.getNextMetastoreEvents(10)); + assertFalse(fixture.searchMatch("orders", 5).isEmpty()); + assertEquals(0, handler.getNextMetastoreEvents(10)); + } + } + + @Test + public void parseFailureFallsBackToSingleEventApply() throws Exception { + Configuration conf = eventHandlerConf(); + MessageBuilder messageBuilder = MessageBuilder.getInstance(); + Table orders = InMemorySearchFixture.table("hive", "sales", "orders", "sales orders"); + NotificationEvent good = createEvent(401L, orders, messageBuilder); + NotificationEvent unparsableBatchMember = event( + 402L, + MessageBuilder.CREATE_TABLE_EVENT, + "{invalid-create-table-message", + "hive", + "sales", + "broken"); + + IMetaStoreClient client = mock(IMetaStoreClient.class); + when(client.getNextNotification(any(NotificationEventRequest.class), eq(true), eq(null))) + .thenReturn(response(good, unparsableBatchMember)); + + try (InMemorySearchFixture fixture = InMemorySearchFixture.create()) { + AtomicBoolean sawBatchFailure = new AtomicBoolean(); + MetastoreEventHandler handler = new MetastoreEventHandler(conf, client); + handler.addListeners(new MetastoreEventListener() { + @Override + public void notifyIndexTask(IndexTask task) throws IOException { + if (task.tablesToAdd.size() > 1) { + sawBatchFailure.set(true); + } + try { + fixture.mutations().apply(task); + fixture.commit(task.lastEventId); + } catch (Exception e) { + throw new IOException(e); + } + } + }); + + assertThrows(IndexNotHealthyException.class, () -> handler.getNextMetastoreEvents(10)); + assertFalse(sawBatchFailure.get()); + assertFalse(fixture.searchMatch("orders", 5).isEmpty()); + } + } + + private static Configuration eventHandlerConf() { + Configuration conf = new Configuration(false); + MetastoreConf.setVar(conf, MetastoreConf.ConfVars.EVENT_MESSAGE_FACTORY, + JSONMessageEncoder.class.getName()); + conf.setLong(IndexConfig.EVENT_FAILURE_BACKOFF_MS, 0L); + conf.setInt(IndexConfig.EVENT_BATCH_MAX_FAILURES, 3); + return conf; + } + + private static NotificationEvent createEvent(long id, Table table, MessageBuilder builder) { + return event( + id, + MessageBuilder.CREATE_TABLE_EVENT, + builder.buildCreateTableMessage(table, null).toString(), + table.getCatName(), + table.getDbName(), + table.getTableName()); + } + + private static NotificationEvent event( + long id, String type, String message, String cat, String db, String table) { + NotificationEvent event = new NotificationEvent(id, 0, type, message); + event.setCatName(cat); + event.setDbName(db); + event.setTableName(table); + return event; + } + + private static NotificationEventResponse response(NotificationEvent... events) { + NotificationEventResponse response = new NotificationEventResponse(); + response.setEvents(List.of(events)); + return response; + } + + private static NotificationEventResponse emptyResponse() { + NotificationEventResponse response = new NotificationEventResponse(); + response.setEvents(new ArrayList<>()); + return response; + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestMetastoreEventHandlerIntegration.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestMetastoreEventHandlerIntegration.java new file mode 100644 index 000000000000..6c606d0eea97 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestMetastoreEventHandlerIntegration.java @@ -0,0 +1,232 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.metastore; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.metastore.IMetaStoreClient; +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hadoop.hive.metastore.conf.MetastoreConf; +import org.apache.hadoop.hive.metastore.messaging.json.JSONMessageEncoder; +import org.apache.hadoop.hive.metastore.api.NotificationEvent; +import org.apache.hadoop.hive.metastore.api.NotificationEventRequest; +import org.apache.hadoop.hive.metastore.api.NotificationEventResponse; +import org.apache.hadoop.hive.metastore.api.Table; +import org.apache.hadoop.hive.metastore.messaging.MessageBuilder; +import org.apache.hive.search.config.IndexConfig; +import org.apache.hive.search.search.InMemorySearchFixture; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@Category(MetastoreUnitTest.class) +public class TestMetastoreEventHandlerIntegration { + + @Test + public void pollsNotificationsAndUpdatesIndex() throws Exception { + Configuration conf = new Configuration(false); + MetastoreConf.setVar(conf, MetastoreConf.ConfVars.EVENT_MESSAGE_FACTORY, + JSONMessageEncoder.class.getName()); + conf.setLong(IndexConfig.EVENT_FAILURE_BACKOFF_MS, 0L); + conf.setInt(IndexConfig.EVENT_BATCH_MAX_FAILURES, 1); + + MessageBuilder messageBuilder = MessageBuilder.getInstance(); + Table salesOrders = InMemorySearchFixture.table("hive", "sales", "orders", "daily sales orders"); + String createMessage = messageBuilder.buildCreateTableMessage(salesOrders, null).toString(); + + NotificationEvent createEvent = event(101L, MessageBuilder.CREATE_TABLE_EVENT, createMessage, + "hive", "sales", "orders"); + NotificationEvent dropEvent = event( + 102L, + MessageBuilder.DROP_TABLE_EVENT, + messageBuilder.buildDropTableMessage(salesOrders).toString(), + "hive", + "sales", + "orders"); + + IMetaStoreClient client = mock(IMetaStoreClient.class); + when(client.getNextNotification(any(NotificationEventRequest.class), eq(true), eq(null))) + .thenReturn(response(createEvent)) + .thenReturn(response(dropEvent)) + .thenReturn(emptyResponse()); + + try (InMemorySearchFixture fixture = InMemorySearchFixture.create()) { + AtomicReference lastTask = new AtomicReference<>(); + MetastoreEventHandler handler = new MetastoreEventHandler(conf, client); + handler.addListeners(new MetastoreEventListener() { + @Override + public void notifyIndexTask(IndexTask task) throws IOException { + lastTask.set(task); + try { + fixture.mutations().apply(task); + fixture.commit(task.lastEventId); + } catch (Exception e) { + throw new IOException(e); + } + } + }); + + assertEquals(1, handler.getNextMetastoreEvents(10)); + assertEquals(1, lastTask.get().tablesToAdd.size()); + assertFalse(fixture.searchMatch("sales", 5).isEmpty()); + + assertEquals(1, handler.getNextMetastoreEvents(10)); + assertEquals(1, lastTask.get().tablesToDrop.size()); + assertTrue(fixture.searchMatch("sales", 5).isEmpty()); + } + } + + @Test + public void skipsAlterWhenIndexedFieldsUnchanged() throws Exception { + Configuration conf = new Configuration(false); + MetastoreConf.setVar(conf, MetastoreConf.ConfVars.EVENT_MESSAGE_FACTORY, + JSONMessageEncoder.class.getName()); + + MessageBuilder messageBuilder = MessageBuilder.getInstance(); + Table before = InMemorySearchFixture.table("hive", "sales", "orders", "daily sales orders"); + Table after = new Table(before); + after.setParameters(new java.util.HashMap<>(before.getParameters())); + after.getParameters().put("transient_lastDdlTime", "999"); + String alterMessage = messageBuilder.buildAlterTableMessage(before, after, false, -1L).toString(); + + NotificationEvent create = event( + 301L, + MessageBuilder.CREATE_TABLE_EVENT, + messageBuilder.buildCreateTableMessage(before, null).toString(), + "hive", + "sales", + "orders"); + NotificationEvent alter = event( + 302L, + MessageBuilder.ALTER_TABLE_EVENT, + alterMessage, + "hive", + "sales", + "orders"); + + IMetaStoreClient client = mock(IMetaStoreClient.class); + when(client.getNextNotification(any(NotificationEventRequest.class), eq(true), eq(null))) + .thenReturn(response(create)) + .thenReturn(response(alter)) + .thenReturn(emptyResponse()); + + try (InMemorySearchFixture fixture = InMemorySearchFixture.create()) { + AtomicReference lastTask = new AtomicReference<>(); + MetastoreEventHandler handler = new MetastoreEventHandler(conf, client); + handler.addListeners(new MetastoreEventListener() { + @Override + public void notifyIndexTask(IndexTask task) throws IOException { + lastTask.set(task); + try { + fixture.mutations().apply(task); + fixture.commit(task.lastEventId); + } catch (Exception e) { + throw new IOException(e); + } + } + }); + + assertEquals(1, handler.getNextMetastoreEvents(10)); + assertEquals(1, lastTask.get().tablesToAdd.size()); + assertFalse(fixture.searchMatch("sales", 5).isEmpty()); + + assertEquals(1, handler.getNextMetastoreEvents(10)); + assertTrue(lastTask.get().tablesToAdd.isEmpty()); + assertTrue(lastTask.get().tablesToDrop.isEmpty()); + assertFalse(fixture.searchMatch("sales", 5).isEmpty()); + } + } + + @Test + public void coalesceCreateThenDropInSameBatch() throws Exception { + Configuration conf = new Configuration(false); + MetastoreConf.setVar(conf, MetastoreConf.ConfVars.EVENT_MESSAGE_FACTORY, + JSONMessageEncoder.class.getName()); + MessageBuilder messageBuilder = MessageBuilder.getInstance(); + Table table = InMemorySearchFixture.table("hive", "sales", "orders", "sales table"); + NotificationEvent create = event( + 201L, + MessageBuilder.CREATE_TABLE_EVENT, + messageBuilder.buildCreateTableMessage(table, null).toString(), + "hive", + "sales", + "orders"); + NotificationEvent drop = event( + 202L, + MessageBuilder.DROP_TABLE_EVENT, + messageBuilder.buildDropTableMessage(table).toString(), + "hive", + "sales", + "orders"); + + IMetaStoreClient client = mock(IMetaStoreClient.class); + when(client.getNextNotification(any(NotificationEventRequest.class), eq(true), eq(null))) + .thenReturn(response(create, drop)) + .thenReturn(emptyResponse()); + + try (InMemorySearchFixture fixture = InMemorySearchFixture.create()) { + MetastoreEventHandler handler = new MetastoreEventHandler(conf, client); + handler.addListeners(new MetastoreEventListener() { + @Override + public void notifyIndexTask(IndexTask task) throws IOException { + try { + fixture.mutations().apply(task); + fixture.commit(task.lastEventId); + } catch (Exception e) { + throw new IOException(e); + } + } + }); + + assertEquals(2, handler.getNextMetastoreEvents(10)); + assertTrue(fixture.searchMatch("sales", 5).isEmpty()); + } + } + + private static NotificationEvent event( + long id, String type, String message, String cat, String db, String table) { + NotificationEvent event = new NotificationEvent(id, 0, type, message); + event.setCatName(cat); + event.setDbName(db); + event.setTableName(table); + return event; + } + + private static NotificationEventResponse response(NotificationEvent... events) { + NotificationEventResponse response = new NotificationEventResponse(); + response.setEvents(List.of(events)); + return response; + } + + private static NotificationEventResponse emptyResponse() { + NotificationEventResponse response = new NotificationEventResponse(); + response.setEvents(new ArrayList<>()); + return response; + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestMetastoreIndexer.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestMetastoreIndexer.java new file mode 100644 index 000000000000..0dbcaf06f2ab --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestMetastoreIndexer.java @@ -0,0 +1,196 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.metastore; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.metastore.IMetaStoreClient; +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hadoop.hive.metastore.api.Table; +import org.apache.hadoop.hive.metastore.conf.MetastoreConf; +import org.apache.hadoop.hive.metastore.messaging.json.JSONMessageEncoder; +import org.apache.hive.search.config.IndexConfig; +import org.apache.hive.search.config.IndexStoreConfig; +import org.apache.hive.search.config.InferenceConfig; +import org.apache.hive.search.config.SearchConfig; +import org.apache.hive.search.index.Indexer; +import org.apache.hive.search.index.IndexManager; +import org.apache.hive.search.index.store.LocalStateClient; +import org.apache.hive.search.inference.EmbedModelRegistry; +import org.apache.hive.search.mapping.IndexMapping; +import org.apache.hive.search.mapping.TableDocument; +import org.apache.hive.search.search.InMemorySearchFixture; +import org.apache.hive.search.testutil.InMemoryIndexStateClient; +import org.apache.hive.search.testutil.MetastoreBootstrapMocks; +import org.apache.hive.search.testutil.StubEmbedModel; +import org.apache.hive.search.testutil.TestLeaderElection; +import org.apache.lucene.store.ByteBuffersDirectory; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.util.List; +import java.util.Map; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; + +@Category(MetastoreUnitTest.class) +public class TestMetastoreIndexer { + + private static final String ELECTION = "metastore-indexer-test"; + + @Test + public void leaderBootstrapsEmptyIndex() throws Exception { + Configuration conf = indexerConf(true); + Table orders = InMemorySearchFixture.table("hive", "sales", "orders", "sales orders"); + Table customers = InMemorySearchFixture.table("hive", "sales", "customers", "sales customers"); + IMetaStoreClient client = mock(IMetaStoreClient.class); + MetastoreBootstrapMocks.stubCurrentNotificationId(client, 500L); + MetastoreBootstrapMocks.stubCatchUp(client); + MetastoreBootstrapMocks.stubBootstrapCatalog(client, orders, customers); + + try (IndexerFixture fixture = IndexerFixture.create(conf, null)) { + try (MetastoreIndexer metastoreIndexer = + new MetastoreIndexer(conf, fixture.indexManager, fixture.indexer, client)) { + assertEquals(500L, fixture.indexManager.readLocalManifest().orElseThrow().lastEventId()); + assertEquals(2, fixture.indexer.writer().getDocStats().numDocs); + } + } + } + + @Test + public void reusesValidLocalIndexWithoutRebuild() throws Exception { + Configuration conf = indexerConf(true); + Table orders = InMemorySearchFixture.table("hive", "sales", "orders", "sales orders"); + IMetaStoreClient client = mock(IMetaStoreClient.class); + MetastoreBootstrapMocks.stubCatchUp(client); + + try (IndexerFixture fixture = IndexerFixture.create(conf, null)) { + fixture.indexer.initialize(); + List docs = List.of( + MetastoreTableMapper.fromTable(orders, fixture.indexManager.mapping())); + fixture.indexer.addDocuments(docs); + fixture.indexer.flush(42L, true); + fixture.indexer.close(); + + fixture.indexer = new Indexer(fixture.indexManager, fixture.modelRegistry); + try (MetastoreIndexer metastoreIndexer = + new MetastoreIndexer(conf, fixture.indexManager, fixture.indexer, client)) { + verify(client, never()).getAllDatabases(); + assertEquals(42L, fixture.indexManager.readLocalManifest().orElseThrow().lastEventId()); + assertEquals(1, fixture.indexer.writer().getDocStats().numDocs); + } + } + } + + @Test + public void followerRestoresFromRemoteBackup() throws Exception { + Configuration conf = indexerConf(false); + Table orders = InMemorySearchFixture.table("hive", "sales", "orders", "sales orders"); + Table customers = InMemorySearchFixture.table("hive", "sales", "customers", "sales customers"); + InMemoryIndexStateClient remote = new InMemoryIndexStateClient(); + + IndexMapping mapping = MetastoreSchemas.defaultHiveTablesMapping( + "test_index", InMemorySearchFixture.MODEL_NAME, conf); + EmbedModelRegistry registry = new EmbedModelRegistry( + Map.of(InMemorySearchFixture.MODEL_NAME, new StubEmbedModel(InMemorySearchFixture.MODEL_NAME))); + + ByteBuffersDirectory leaderDir = new ByteBuffersDirectory(); + IndexManager leaderManager = new IndexManager( + mapping, leaderDir, new LocalStateClient(leaderDir, "test_index"), remote); + Indexer leaderIndexer = new Indexer(leaderManager, registry); + leaderIndexer.initialize(); + leaderIndexer.addDocuments(List.of( + MetastoreTableMapper.fromTable(orders, mapping), + MetastoreTableMapper.fromTable(customers, mapping))); + leaderIndexer.flush(200L, true); + leaderIndexer.syncBackup(); + leaderIndexer.close(); + leaderManager.close(); + + IMetaStoreClient client = mock(IMetaStoreClient.class); + MetastoreBootstrapMocks.stubCatchUp(client); + + try (IndexerFixture fixture = IndexerFixture.create(conf, remote)) { + try (MetastoreIndexer metastoreIndexer = + new MetastoreIndexer(conf, fixture.indexManager, fixture.indexer, client)) { + verify(client, never()).getAllDatabases(); + assertEquals(200L, fixture.indexManager.readLocalManifest().orElseThrow().lastEventId()); + assertEquals(2, fixture.indexer.writer().getDocStats().numDocs); + } + } + } + + private static Configuration indexerConf(boolean leader) { + Configuration conf = new Configuration(false); + conf.setBoolean(IndexStoreConfig.MEMORY, true); + conf.set(IndexConfig.INDEX_NAME, "test_index"); + conf.set(InferenceConfig.MODEL_NAME, InMemorySearchFixture.MODEL_NAME); + conf.setInt(IndexConfig.BOOTSTRAP_FETCH_THREADS, 1); + conf.setInt(IndexConfig.BOOTSTRAP_QUEUE_DEPTH, 4); + conf.setLong(IndexConfig.BOOTSTRAP_PROGRESS_INTERVAL_MS, Long.MAX_VALUE); + conf.setLong(IndexConfig.EVENT_FAILURE_BACKOFF_MS, 0L); + conf.setLong(IndexConfig.FLUSH_INTERVAL_MS, 60_000L); + conf.setLong(IndexConfig.INDEX_SYNC_INTERVAL, 60_000L); + conf.setInt(SearchConfig.BAYESIAN_SAMPLES, 5); + conf.setInt(SearchConfig.BAYESIAN_TOKENS_PER_QUERY, 2); + conf.setLong(SearchConfig.BAYESIAN_SEED, 1L); + MetastoreConf.setVar(conf, MetastoreConf.ConfVars.EVENT_MESSAGE_FACTORY, + JSONMessageEncoder.class.getName()); + TestLeaderElection election = new TestLeaderElection(leader); + MetastoreCluster.injectElection(conf, ELECTION, election, "mutex"); + return conf; + } + + private static final class IndexerFixture implements AutoCloseable { + IndexManager indexManager; + Indexer indexer; + final EmbedModelRegistry modelRegistry; + + private IndexerFixture( + IndexManager indexManager, Indexer indexer, EmbedModelRegistry modelRegistry) { + this.indexManager = indexManager; + this.indexer = indexer; + this.modelRegistry = modelRegistry; + } + + static IndexerFixture create(Configuration conf, InMemoryIndexStateClient remote) + throws Exception { + IndexMapping mapping = MetastoreSchemas.defaultHiveTablesMapping( + "test_index", InMemorySearchFixture.MODEL_NAME, conf); + EmbedModelRegistry registry = new EmbedModelRegistry( + Map.of(InMemorySearchFixture.MODEL_NAME, new StubEmbedModel(InMemorySearchFixture.MODEL_NAME))); + ByteBuffersDirectory directory = new ByteBuffersDirectory(); + LocalStateClient local = new LocalStateClient(directory, "test_index"); + IndexManager indexManager = remote == null + ? new IndexManager(mapping, directory, local, null) + : new IndexManager(mapping, directory, local, remote); + Indexer indexer = new Indexer(indexManager, registry); + return new IndexerFixture(indexManager, indexer, registry); + } + + @Override + public void close() throws Exception { + indexer.close(); + indexManager.close(); + modelRegistry.close(); + } + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestMetastoreTableMapper.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestMetastoreTableMapper.java new file mode 100644 index 000000000000..042df380dfa8 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestMetastoreTableMapper.java @@ -0,0 +1,385 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.metastore; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.common.TableName; +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hadoop.hive.metastore.api.FieldSchema; +import org.apache.hadoop.hive.metastore.api.StorageDescriptor; +import org.apache.hadoop.hive.metastore.api.Table; +import org.apache.hive.search.config.IndexStoreConfig; +import org.apache.hive.search.index.IndexManager; +import org.apache.hive.search.index.Indexer; +import org.apache.hive.search.inference.EmbedModelRegistry; +import org.apache.hive.search.mapping.IndexMapping; +import org.apache.hive.search.mapping.TableDocument; +import org.apache.hive.search.mapping.field.Field; +import org.apache.hive.search.mapping.field.IdField; +import org.apache.hive.search.mapping.field.TextField; +import org.apache.lucene.document.Document; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.IndexableField; +import org.apache.lucene.index.MultiTerms; +import org.apache.lucene.store.ByteBuffersDirectory; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +@Category(MetastoreUnitTest.class) +public class TestMetastoreTableMapper { + + @Test + public void tableIdIncludesCatalogDbAndTable() { + assertEquals("hive.default.orders", + MetastoreTableMapper.tableId("hive", "default", "orders")); + assertEquals("hive.default.orders", + MetastoreTableMapper.tableId(TableName.fromString("default.orders", "hive", "default"))); + } + + @Test + public void fromTableBuildsSearchTextAndStoredFields() throws Exception { + Configuration conf = new Configuration(false); + IndexMapping mapping = MetastoreSchemas.defaultHiveTablesMapping("hive_tables", "bge-small", conf); + + Table table = new Table(); + table.setCatName("hive"); + table.setDbName("sales"); + table.setTableName("orders"); + table.setOwner("alice"); + table.setTableType("MANAGED_TABLE"); + table.setSd(new StorageDescriptor()); + table.getSd().setLocation("hdfs://warehouse/orders"); + table.getSd().setCols(List.of( + new FieldSchema("id", "bigint", "order id"), + new FieldSchema("amount", "double", null))); + Map params = new HashMap<>(); + params.put("comment", "daily orders"); + table.setParameters(params); + + TableDocument document = MetastoreTableMapper.fromTable(table, mapping); + document = withSearchTextEmbedding(document, mapping, new float[] {0.1f, 0.2f, 0.3f}); + assertEquals("hive.sales.orders", document.idField().value()); + + List luceneDocs = document.toDocuments(); + assertEquals(1, luceneDocs.size()); + Document luceneDoc = luceneDocs.get(0); + assertTrue(luceneDoc.getFields().size() >= 10); + assertTrue(luceneDoc.get("_id").contains("hive.sales.orders")); + assertEquals("sales", luceneDoc.get(MetastoreTableMapper.FIELD_DB)); + assertEquals("orders", luceneDoc.get(MetastoreTableMapper.FIELD_TABLE)); + assertTrue(hasIndexedField(luceneDoc, MetastoreTableMapper.FIELD_TABLE + ".filter")); + String searchText = luceneDoc.get(MetastoreTableMapper.FIELD_SEARCH_TEXT); + assertEquals( + "table: orders; comment: daily orders; column id: order id; column amount:", + searchText); + assertFalse(searchText.contains("hdfs://")); + assertFalse(searchText.contains("MANAGED_TABLE")); + assertFalse(searchText.contains("alice")); + assertEquals("id bigint order id; amount double", luceneDoc.get(MetastoreTableMapper.FIELD_COLUMNS)); + assertEquals("id amount", luceneDoc.get(MetastoreTableMapper.FIELD_COLUMN_NAMES)); + assertEquals("order id", luceneDoc.get(MetastoreTableMapper.FIELD_COLUMN_COMMENTS)); + } + + @Test + public void columnSearchFieldsSplitNamesAndComments() throws Exception { + Configuration conf = new Configuration(false); + IndexMapping mapping = MetastoreSchemas.defaultHiveTablesMapping("hive_tables", "bge-small", conf); + + Table table = new Table(); + table.setCatName("hive"); + table.setDbName("sales"); + table.setTableName("orders"); + table.setSd(new StorageDescriptor()); + table.getSd().setCols(List.of( + new FieldSchema("id", "bigint", "order id"), + new FieldSchema("amount", "double", null), + new FieldSchema("status", "string", "fulfillment status"))); + + TableDocument document = MetastoreTableMapper.fromTable(table, mapping); + assertEquals("id amount status", fieldValue(document, MetastoreTableMapper.FIELD_COLUMN_NAMES)); + assertEquals("order id; fulfillment status", + fieldValue(document, MetastoreTableMapper.FIELD_COLUMN_COMMENTS)); + } + + @Test + public void searchTextIncludesAllColumns() throws Exception { + Configuration conf = new Configuration(false); + IndexMapping mapping = MetastoreSchemas.defaultHiveTablesMapping("hive_tables", "bge-small", conf); + + Table table = new Table(); + table.setCatName("hive"); + table.setDbName("sales"); + table.setTableName("orders"); + table.setSd(new StorageDescriptor()); + table.getSd().setCols(List.of( + new FieldSchema("id", "bigint", "order id"), + new FieldSchema("amount", "double", null), + new FieldSchema("status", "string", "fulfillment status"))); + + TableDocument document = MetastoreTableMapper.fromTable(table, mapping); + assertEquals( + "table: orders; column id: order id; column amount:; column status: fulfillment status", + fieldValue(document, MetastoreTableMapper.FIELD_SEARCH_TEXT)); + assertTrue(fieldValue(document, MetastoreTableMapper.FIELD_COLUMNS).contains("amount double")); + } + + @Test + public void searchTextIncludesAllCommentedColumnsForWideTables() throws Exception { + Configuration conf = new Configuration(false); + IndexMapping mapping = MetastoreSchemas.defaultHiveTablesMapping("hive_tables", "bge-small", conf); + + Table table = new Table(); + table.setCatName("hive"); + table.setDbName("wide"); + table.setTableName("events"); + table.setSd(new StorageDescriptor()); + List cols = new ArrayList<>(); + int columnCount = 105; + for (int i = 0; i < columnCount; i++) { + cols.add(new FieldSchema("col" + i, "string", "comment " + i)); + } + table.getSd().setCols(cols); + + TableDocument document = MetastoreTableMapper.fromTable(table, mapping); + String searchText = fieldValue(document, MetastoreTableMapper.FIELD_SEARCH_TEXT); + String columnComments = fieldValue(document, MetastoreTableMapper.FIELD_COLUMN_COMMENTS); + String storedColumns = fieldValue(document, MetastoreTableMapper.FIELD_COLUMNS); + + assertTrue(searchText.contains("column col0: comment 0")); + assertTrue(searchText.contains( + "column col" + (columnCount - 1) + ": comment " + (columnCount - 1))); + assertFalse(searchText.contains("(+")); + assertTrue(columnComments.contains("comment " + (columnCount - 1))); + assertFalse(columnComments.contains("(+")); + assertTrue(storedColumns.contains("col" + (columnCount - 1) + " string comment " + (columnCount - 1))); + assertFalse(storedColumns.contains("(+")); + } + + @Test + public void searchTextUsesStructuredLabelsWithoutComment() throws Exception { + Configuration conf = new Configuration(false); + IndexMapping mapping = MetastoreSchemas.defaultHiveTablesMapping("hive_tables", "bge-small", conf); + + Table table = new Table(); + table.setCatName("hive"); + table.setDbName("sales"); + table.setTableName("Events"); + table.setSd(new StorageDescriptor()); + table.getSd().setCols(List.of(new FieldSchema("id", "bigint", null))); + + TableDocument document = MetastoreTableMapper.fromTable(table, mapping); + assertEquals("table: events; column id:", fieldValue(document, MetastoreTableMapper.FIELD_SEARCH_TEXT)); + } + + @Test + public void embedDocumentsPreservesLexicalFields() throws Exception { + Configuration conf = new Configuration(false); + conf.setBoolean(IndexStoreConfig.MEMORY, true); + conf.set(org.apache.hive.search.config.InferenceConfig.MODEL_NAME, "stub-model"); + IndexMapping mapping = MetastoreSchemas.defaultHiveTablesMapping("hive_tables", "stub-model", conf); + IndexManager indexManager = IndexManager.open(mapping, conf); + EmbedModelRegistry registry = new EmbedModelRegistry( + java.util.Map.of("stub-model", new org.apache.hive.search.testutil.StubEmbedModel("stub-model"))); + Indexer indexer = new Indexer(indexManager, registry); + indexer.initialize(); + + Table table = sampleTable(); + Map params = new HashMap<>(); + params.put("comment", "daily sales orders"); + table.setParameters(params); + + TableDocument doc = MetastoreTableMapper.fromTable(table, mapping); + TableDocument embedded = indexer.embedDocuments(java.util.List.of(doc)).get(0); + java.util.Set fieldNames = new java.util.HashSet<>(); + for (Field field : embedded.fields()) { + if (field instanceof TextField textField) { + fieldNames.add(textField.name()); + } + } + assertTrue(fieldNames.contains(MetastoreTableMapper.FIELD_COMMENT)); + assertTrue(fieldNames.contains(MetastoreTableMapper.FIELD_SEARCH_TEXT)); + + ByteBuffersDirectory directory = new ByteBuffersDirectory(); + try (IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig(mapping.analyzer()))) { + writer.addDocuments(embedded.toDocuments()); + } + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertNotNull(MultiTerms.getTerms(reader, MetastoreTableMapper.FIELD_COMMENT)); + } + indexer.close(); + indexManager.close(); + registry.close(); + } + + @Test + public void lexicalFieldsAreIndexedForKeywordSearch() throws Exception { + Configuration conf = new Configuration(false); + IndexMapping mapping = MetastoreSchemas.defaultHiveTablesMapping("hive_tables", "bge-small", conf); + + Table table = new Table(); + table.setCatName("hive"); + table.setDbName("sales"); + table.setTableName("orders"); + table.setSd(new StorageDescriptor()); + Map params = new HashMap<>(); + params.put("comment", "daily sales orders"); + table.setParameters(params); + + TableDocument document = MetastoreTableMapper.fromTable(table, mapping); + document = withSearchTextEmbedding(document, mapping, new float[] {0.1f, 0.2f, 0.3f}); + ByteBuffersDirectory directory = new ByteBuffersDirectory(); + try (IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig(mapping.analyzer()))) { + writer.addDocuments(document.toDocuments()); + } + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertNotNull(MultiTerms.getTerms(reader, MetastoreTableMapper.FIELD_COMMENT)); + assertNotNull(MultiTerms.getTerms(reader, MetastoreTableMapper.FIELD_TABLE)); + assertNotNull(MultiTerms.getTerms(reader, MetastoreTableMapper.FIELD_TABLE + ".filter")); + } + } + + @Test + public void semanticFieldRequiresEmbedding() throws Exception { + Configuration conf = new Configuration(false); + IndexMapping mapping = MetastoreSchemas.defaultHiveTablesMapping("hive_tables", "bge-small", conf); + TableDocument document = MetastoreTableMapper.fromTable(sampleTable(), mapping); + document.appendField(new TextField(MetastoreTableMapper.FIELD_SEARCH_TEXT, "sales data")); + try { + document.toDocuments(); + org.junit.Assert.fail("expected semantic field without embedding to fail"); + } catch (org.apache.hive.search.exception.IndexException expected) { + assertTrue(expected.getMessage().contains("requires embedding")); + } + } + + @Test + public void hasIndexedFieldsChangedDetectsIndexedAlterations() { + Table before = sampleIndexedTable(); + Table after = copyTable(before); + assertFalse(MetastoreTableMapper.hasIndexedFieldsChanged(before, after)); + + after.getParameters().put("transient_lastDdlTime", "123"); + assertFalse(MetastoreTableMapper.hasIndexedFieldsChanged(before, after)); + + after = copyTable(before); + after.getParameters().put("comment", "updated comment"); + assertTrue(MetastoreTableMapper.hasIndexedFieldsChanged(before, after)); + + after = copyTable(before); + after.setOwner("bob"); + assertTrue(MetastoreTableMapper.hasIndexedFieldsChanged(before, after)); + + after = copyTable(before); + after.setTableName("invoices"); + assertTrue(MetastoreTableMapper.hasIndexedFieldsChanged(before, after)); + + after = copyTable(before); + after.getSd().getCols().get(0).setComment("new column comment"); + assertTrue(MetastoreTableMapper.hasIndexedFieldsChanged(before, after)); + + after = copyTable(before); + after.getSd().getCols().get(1).setType("int"); + assertTrue(MetastoreTableMapper.hasIndexedFieldsChanged(before, after)); + } + + private static Table sampleIndexedTable() { + Table table = new Table(); + table.setCatName("hive"); + table.setDbName("sales"); + table.setTableName("orders"); + table.setOwner("alice"); + table.setTableType("MANAGED_TABLE"); + table.setSd(new StorageDescriptor()); + table.getSd().setLocation("hdfs://warehouse/orders"); + table.getSd().setCols(List.of( + new FieldSchema("id", "bigint", "order id"), + new FieldSchema("amount", "double", null))); + Map params = new HashMap<>(); + params.put("comment", "daily orders"); + table.setParameters(params); + return table; + } + + private static Table copyTable(Table source) { + Table copy = new Table(source); + copy.setSd(new StorageDescriptor(source.getSd())); + List cols = new ArrayList<>(); + for (FieldSchema column : source.getSd().getCols()) { + cols.add(new FieldSchema(column)); + } + copy.getSd().setCols(cols); + copy.setParameters(new HashMap<>(source.getParameters())); + return copy; + } + + private static Table sampleTable() { + Table table = new Table(); + table.setCatName("hive"); + table.setDbName("default"); + table.setTableName("t"); + table.setSd(new StorageDescriptor()); + return table; + } + + private static TableDocument withSearchTextEmbedding( + TableDocument document, IndexMapping mapping, float[] embedding) { + java.util.List fields = new java.util.ArrayList<>(); + for (Field field : document.fields()) { + if (field instanceof IdField) { + continue; + } + if (field instanceof TextField textField + && MetastoreTableMapper.FIELD_SEARCH_TEXT.equals(textField.name())) { + fields.add(textField.withEmbedding(embedding)); + } else { + fields.add(field); + } + } + return new TableDocument(document.idField(), fields, mapping); + } + + private static boolean hasIndexedField(Document document, String fieldName) { + for (IndexableField field : document.getFields()) { + if (fieldName.equals(field.name())) { + return true; + } + } + return false; + } + + private static String fieldValue(TableDocument document, String fieldName) { + for (Field field : document.fields()) { + if (field instanceof TextField textField && fieldName.equals(textField.name())) { + return textField.value(); + } + } + throw new AssertionError("missing field: " + fieldName); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestRealMetastoreSearchFailureRecovery.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestRealMetastoreSearchFailureRecovery.java new file mode 100644 index 000000000000..9ae3324c3876 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestRealMetastoreSearchFailureRecovery.java @@ -0,0 +1,113 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.metastore; + +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hive.search.exception.IndexNotHealthyException; +import org.apache.hive.search.search.TableSearchResult; +import org.apache.hive.search.testutil.InMemoryIndexStateClient; +import org.apache.hive.search.testutil.RealMetastoreServer; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +@Category(MetastoreUnitTest.class) +public class TestRealMetastoreSearchFailureRecovery { + + private static RealMetastoreServer metastore; + + @BeforeClass + public static void startMetastore() throws Exception { + metastore = RealMetastoreServer.start(); + } + + @AfterClass + public static void stopMetastore() throws Exception { + if (metastore != null) { + metastore.close(); + } + } + + @Test + public void restoresIndexFromRemoteBackupAfterLocalLoss() throws Exception { + String db = RealMetastoreSearchSession.uniqueDbName("restore"); + metastore.createDatabase(db); + metastore.createTable(db, "orders", "daily sales orders"); + + InMemoryIndexStateClient remote = RealMetastoreSearchSession.newSharedRemoteBackup(); + try (RealMetastoreSearchSession leader = RealMetastoreSearchSession.open(metastore, remote)) { + leader.waitUntilSearchable("sales", 5); + leader.syncBackup(); + assertTrue(remote.readManifest().isPresent()); + } + + try (RealMetastoreSearchSession restored = RealMetastoreSearchSession.open(metastore, remote)) { + restored.refreshSearcher(); + TableSearchResult result = restored.searchMatch("sales", 5); + assertFalse(result.hits().isEmpty()); + assertTrue(result.hits().stream().anyMatch(hit -> hit.table().getTable().contains("orders"))); + } + } + + @Test + public void recoversIncrementalUpdatesAfterRestore() throws Exception { + String db = RealMetastoreSearchSession.uniqueDbName("restore_inc"); + metastore.createDatabase(db); + metastore.createTable(db, "orders", "daily sales orders"); + + InMemoryIndexStateClient remote = RealMetastoreSearchSession.newSharedRemoteBackup(); + try (RealMetastoreSearchSession leader = RealMetastoreSearchSession.open(metastore, remote)) { + leader.waitUntilSearchable("sales", 5); + leader.syncBackup(); + } + + try (RealMetastoreSearchSession restored = RealMetastoreSearchSession.open(metastore, remote)) { + restored.waitUntilSearchable("sales", 5); + metastore.createTable(db, "customers", "customer master records"); + restored.waitUntilSearchable("customer", 5); + assertEquals(2, restored.searchMatch("master", 10).hits().size()); + } + } + + @Test + public void skipPoisonEventAllowsIndexingToContinue() throws Exception { + String db = RealMetastoreSearchSession.uniqueDbName("poison"); + metastore.createDatabase(db); + metastore.createTable(db, "orders", "daily sales orders"); + + try (RealMetastoreSearchSession session = RealMetastoreSearchSession.open(metastore)) { + session.waitUntilSearchable("sales", 5); + session.indexManager().notifyIndexState(false, + new IndexNotHealthyException("simulated notification failure")); + assertThrows(IndexNotHealthyException.class, session.indexManager()::checkIndexState); + + session.indexManager().notifyIndexState(true); + session.indexManager().checkIndexState(); + + metastore.createTable(db, "customers", "customer master records"); + session.waitUntilSearchable("customer", 5); + assertFalse(session.searchMatch("customer", 5).hits().isEmpty()); + } + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestRealMetastoreSearchHybrid.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestRealMetastoreSearchHybrid.java new file mode 100644 index 000000000000..0926de8db344 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestRealMetastoreSearchHybrid.java @@ -0,0 +1,98 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.metastore; + +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hive.search.search.TableSearchResult; +import org.apache.hive.search.testutil.RealMetastoreServer; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +@Category(MetastoreUnitTest.class) +public class TestRealMetastoreSearchHybrid { + + private static RealMetastoreServer metastore; + + @BeforeClass + public static void startMetastore() throws Exception { + metastore = RealMetastoreServer.start(); + } + + @AfterClass + public static void stopMetastore() throws Exception { + if (metastore != null) { + metastore.close(); + } + } + + @Test + public void keywordSearchFindsMatchingTable() throws Exception { + String db = RealMetastoreSearchSession.uniqueDbName("keyword"); + metastore.createDatabase(db); + metastore.createTable(db, "orders", "daily sales orders"); + metastore.createTable(db, "inventory", "spare parts catalog"); + + try (RealMetastoreSearchSession session = RealMetastoreSearchSession.open(metastore)) { + session.waitUntilSearchable("sales", 5); + TableSearchResult result = session.searchMatch("sales", 5); + assertFalse(result.hits().isEmpty()); + assertTrue(result.hits().stream().anyMatch(hit -> hit.table().getTable().contains("orders"))); + assertTrue(session.searchMatch("parts", 5).hits().stream() + .anyMatch(hit -> hit.table().getTable().contains("inventory"))); + } + } + + @Test + public void hybridSearchFindsSemanticallyNamedTable() throws Exception { + String db = RealMetastoreSearchSession.uniqueDbName("hybrid"); + metastore.createDatabase(db); + metastore.createTable(db, "orders", "revenue analytics dashboard"); + metastore.createTable(db, "parts", "spare parts catalog"); + + try (RealMetastoreSearchSession session = RealMetastoreSearchSession.open(metastore)) { + session.waitUntilSearchable("revenue", 5); + TableSearchResult searchResult = session.searchHybrid("revenue analytics", 5); + assertFalse(searchResult.hits().isEmpty()); + assertTrue(searchResult.hits().stream().anyMatch(hit -> hit.table().getTable().contains("orders"))); + } + } + + @Test + public void hybridSearchRanksKeywordMatchHigherForExactToken() throws Exception { + String db = RealMetastoreSearchSession.uniqueDbName("hybrid_rank"); + metastore.createDatabase(db); + metastore.createTable(db, "orders", "sales orders table"); + metastore.createTable(db, "metrics", "revenue analytics dashboard"); + + try (RealMetastoreSearchSession session = RealMetastoreSearchSession.open(metastore)) { + session.waitUntilSearchable("sales", 5); + TableSearchResult result = session.searchMatch("sales", 5); + assertFalse(result.hits().isEmpty()); + assertTrue(result.hits().get(0).table().getTable().contains("orders")); + + result = session.searchHybrid("revenue analytics", 5); + assertFalse(result.hits().isEmpty()); + assertTrue(result.hits().stream().anyMatch(hit -> hit.table().getTable().contains("metrics"))); + } + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestRealMetastoreSearchIntegration.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestRealMetastoreSearchIntegration.java new file mode 100644 index 000000000000..0c2a467ec7db --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestRealMetastoreSearchIntegration.java @@ -0,0 +1,96 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.metastore; + +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hive.search.testutil.RealMetastoreServer; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.util.List; +import java.util.Map; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +@Category(MetastoreUnitTest.class) +public class TestRealMetastoreSearchIntegration { + + private static RealMetastoreServer metastore; + + @BeforeClass + public static void startMetastore() throws Exception { + metastore = RealMetastoreServer.start(); + } + + @AfterClass + public static void stopMetastore() throws Exception { + if (metastore != null) { + metastore.close(); + } + } + + @Test + public void bootstrapIndexesExistingTables() throws Exception { + String db = RealMetastoreSearchSession.uniqueDbName("bootstrap"); + metastore.createDatabase(db); + metastore.createTable(db, "orders", "daily sales orders"); + metastore.createTable(db, "customers", "customer master data"); + + try (RealMetastoreSearchSession session = RealMetastoreSearchSession.open(metastore)) { + session.refreshSearcher(); + assertFalse(session.searchMatch("orders", 10).hits().isEmpty()); + assertFalse(session.searchMatch("customers", 10).hits().isEmpty()); + } + } + + @Test + public void createAndDropTableUpdatesSearchIndex() throws Exception { + String db = RealMetastoreSearchSession.uniqueDbName("incremental"); + metastore.createDatabase(db); + + try (RealMetastoreSearchSession session = RealMetastoreSearchSession.open(metastore)) { + metastore.createTable(db, "orders", "daily sales orders"); + session.waitUntilSearchable("sales", 5); + assertFalse(session.searchMatch("sales", 5).hits().isEmpty()); + + metastore.dropTable(db, "orders"); + session.waitUntilNotSearchable("sales", 5); + assertTrue(session.searchMatch("sales", 5).hits().isEmpty()); + } + } + + @Test + public void alterTableCommentUpdatesKeywordSearch() throws Exception { + String db = RealMetastoreSearchSession.uniqueDbName("alter"); + metastore.createDatabase(db); + metastore.createTable(db, "orders", "old warehouse comment"); + + try (RealMetastoreSearchSession session = RealMetastoreSearchSession.open(metastore)) { + session.waitUntilSearchable("warehouse", 5); + assertTrue(session.searchMatch("revenue", 5).hits().isEmpty()); + + metastore.dropTable(db, "orders"); + metastore.createTable(db, "orders", "revenue analytics dashboard"); + session.waitUntilSearchable("revenue", 5); + assertFalse(session.searchMatch("revenue", 5).hits().isEmpty()); + } + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/search/InMemorySearchFixture.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/search/InMemorySearchFixture.java new file mode 100644 index 000000000000..2056def2e071 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/search/InMemorySearchFixture.java @@ -0,0 +1,171 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.search; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.metastore.api.StorageDescriptor; +import org.apache.hadoop.hive.metastore.api.Table; +import org.apache.hive.search.config.IndexConfig; +import org.apache.hive.search.config.IndexStoreConfig; +import org.apache.hive.search.config.InferenceConfig; +import org.apache.hive.search.config.SearchConfig; +import org.apache.hive.search.index.Indexer; +import org.apache.hive.search.index.IndexManager; +import org.apache.hive.search.inference.EmbedModelRegistry; +import org.apache.hive.search.mapping.IndexMapping; +import org.apache.hive.search.metastore.MetastoreSchemas; +import org.apache.hive.search.metastore.MetastoreTableMapper; +import org.apache.hive.search.testutil.IndexMutationApplier; +import org.apache.hive.search.testutil.StubEmbedModel; +import org.apache.lucene.search.BayesianScoreEstimator; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.SearcherManager; + +import java.io.IOException; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** In-memory Lucene index wired with a stub embedding model. */ +public final class InMemorySearchFixture implements AutoCloseable { + public static final String MODEL_NAME = "stub-model"; + + private final IndexManager indexManager; + private final Indexer indexer; + private final EmbedModelRegistry modelRegistry; + private final SearchConfig searchConfig; + private final IndexMutationApplier mutations; + private SearcherManager searcherManager; + private BayesianScoreEstimator.Parameters bayesianParameters; + + private InMemorySearchFixture( + IndexManager indexManager, + Indexer indexer, + EmbedModelRegistry modelRegistry, + SearchConfig searchConfig) { + this.indexManager = indexManager; + this.indexer = indexer; + this.modelRegistry = modelRegistry; + this.searchConfig = searchConfig; + this.mutations = new IndexMutationApplier(indexManager, indexer); + } + + public static InMemorySearchFixture create() throws Exception { + Configuration conf = new Configuration(false); + conf.setBoolean(IndexStoreConfig.MEMORY, true); + conf.set(IndexConfig.INDEX_NAME, "test_index"); + conf.set(InferenceConfig.MODEL_NAME, MODEL_NAME); + conf.setInt(SearchConfig.BAYESIAN_SAMPLES, 5); + conf.setInt(SearchConfig.BAYESIAN_TOKENS_PER_QUERY, 2); + conf.setLong(SearchConfig.BAYESIAN_SEED, 1L); + + IndexMapping mapping = MetastoreSchemas.defaultHiveTablesMapping( + "test_index", MODEL_NAME, conf); + IndexManager indexManager = IndexManager.open(mapping, conf); + EmbedModelRegistry registry = + new EmbedModelRegistry(Map.of(MODEL_NAME, new StubEmbedModel(MODEL_NAME))); + Indexer indexer = new Indexer(indexManager, registry); + indexer.initialize(); + return new InMemorySearchFixture(indexManager, indexer, registry, new SearchConfig(conf)); + } + + public IndexMutationApplier mutations() { + return mutations; + } + + public void commit(long eventId) throws IOException { + indexer.flush(eventId, true); + refreshSearcher(); + } + + public void refreshSearcher() throws IOException { + if (searcherManager == null) { + searcherManager = new SearcherManager(indexer.writer(), null); + } else { + searcherManager.maybeRefresh(); + } + if (bayesianParameters == null) { + IndexSearcher searcher = searcherManager.acquire(); + try { + bayesianParameters = BayesianScoreEstimator.estimate( + searcher, + MetastoreTableMapper.FIELD_SEARCH_TEXT, + searchConfig.getBayesianSamples(), + searchConfig.getBayesianTokensPerQuery(), + searchConfig.getBayesianSeed()); + } finally { + searcherManager.release(searcher); + } + } + } + + public List searchMatch(String text, int limit) throws Exception { + return search(new SearchArgs.Match(text), limit); + } + + public List search(SearchArgs args, int limit) throws Exception { + if (searcherManager == null || bayesianParameters == null) { + throw new IllegalStateException("call commit() before searching"); + } + try (SearchInternal searchIO = new SearchInternal( + searcherManager, indexManager, modelRegistry, searchConfig, bayesianParameters)) { + TableSearchResult result = searchIO.search(new SearchQuery( + args, + null, null, limit, + List.of(MetastoreTableMapper.FIELD_TABLE, MetastoreTableMapper.FIELD_COMMENT))); + return result.hits(); + } + } + + public static Table table(String catalog, String db, String name, String comment) { + Table table = new Table(); + table.setCatName(catalog); + table.setDbName(db); + table.setTableName(name); + table.setOwner("owner"); + table.setTableType("MANAGED_TABLE"); + table.setSd(new StorageDescriptor()); + table.getSd().setLocation("hdfs://warehouse/" + db + "/" + name); + Map params = new HashMap<>(); + params.put("comment", comment); + table.setParameters(params); + return table; + } + + public IndexManager indexManager() { + return indexManager; + } + + public Indexer indexer() { + return indexer; + } + + public SearcherManager searcherManager() { + return searcherManager; + } + + @Override + public void close() throws Exception { + if (searcherManager != null) { + searcherManager.close(); + } + indexer.close(); + indexManager.close(); + modelRegistry.close(); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/search/TestHybridSearch.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/search/TestHybridSearch.java new file mode 100644 index 000000000000..29930303b713 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/search/TestHybridSearch.java @@ -0,0 +1,51 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.search; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hive.search.config.SearchConfig; +import org.apache.hive.search.exception.SearchException; +import org.apache.hive.search.metastore.MetastoreSchemas; +import org.apache.hive.search.metastore.MetastoreTableMapper; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +@Category(MetastoreUnitTest.class) +public class TestHybridSearch { + + @Test + public void toFusionRequestUsesDefaultsFromSearchConfig() throws SearchException { + HybridSearch.ResolvedHybridQuery resolved = new HybridSearch.ResolvedHybridQuery( + MetastoreTableMapper.FIELD_SEARCH_TEXT, "sales", null); + SearchConfig searchConfig = new SearchConfig(new Configuration(false)); + SearchInternal.FusionRequest request = + HybridSearch.toFusionRequest(resolved, 20, searchConfig); + + assertEquals(20, request.size()); + assertEquals(2, request.retrievers().size()); + assertEquals(SearchQuery.Mode.MATCH.name(), request.retrievers().get(0).name()); + assertTrue(request.retrievers().get(0).query() instanceof SearchArgs.Match); + assertEquals("sales", ((SearchArgs.Match) request.retrievers().get(0).query()).queryText()); + assertEquals(searchConfig.getHybridMatchWeight(), request.retrievers().get(0).weight(), 0.001f); + assertEquals(searchConfig.getHybridSemanticWeight(), request.retrievers().get(1).weight(), 0.001f); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/search/TestSearchArgs.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/search/TestSearchArgs.java new file mode 100644 index 000000000000..e5f276096b29 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/search/TestSearchArgs.java @@ -0,0 +1,226 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.search; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hive.search.exception.SearchException; +import org.apache.hive.search.mapping.FieldSchema; +import org.apache.hive.search.mapping.IndexMapping; +import org.apache.hive.search.mapping.SearchParams; +import org.apache.hive.search.metastore.MetastoreSchemas; +import org.apache.hive.search.metastore.MetastoreTableMapper; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.util.LinkedHashMap; +import java.util.Map; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; + +@Category(MetastoreUnitTest.class) +public class TestSearchArgs { + + private static IndexMapping hybridMapping() { + Configuration conf = new Configuration(false); + return MetastoreSchemas.defaultHiveTablesMapping("hive_tables", "bge-small", conf); + } + + @Test + public void parseHybridString() throws Exception { + SearchArgs.Hybrid hybrid = + (SearchArgs.Hybrid) SearchArgs.fromBody("sales revenue", SearchQuery.Mode.HYBRID); + assertEquals("sales revenue", hybrid.queryText()); + assertNull(hybrid.field()); + } + + @Test + public void parseHybridFieldAndQueryObject() throws Exception { + Map body = Map.of( + "field", MetastoreTableMapper.FIELD_SEARCH_TEXT, + "query", "orders"); + SearchArgs.Hybrid hybrid = + (SearchArgs.Hybrid) SearchArgs.fromBody(body, SearchQuery.Mode.HYBRID); + assertEquals(MetastoreTableMapper.FIELD_SEARCH_TEXT, hybrid.field()); + assertEquals("orders", hybrid.queryText()); + } + + @Test + public void parseHybridSingleFieldMap() throws Exception { + Map body = Map.of(MetastoreTableMapper.FIELD_SEARCH_TEXT, "inventory"); + SearchArgs.Hybrid hybrid = + (SearchArgs.Hybrid) SearchArgs.fromBody(body, SearchQuery.Mode.HYBRID); + assertEquals("inventory", hybrid.queryText()); + } + + @Test + public void parseHybridCustomWeights() throws Exception { + Map body = Map.of( + "field", MetastoreTableMapper.FIELD_SEARCH_TEXT, + "query", "metrics", + "match_weight", 0.7, + "semantic_weight", 0.3); + SearchArgs.Hybrid hybrid = + (SearchArgs.Hybrid) SearchArgs.fromBody(body, SearchQuery.Mode.HYBRID); + assertEquals(0.3f, hybrid.semanticWeight(), 0.001f); + } + + @Test + public void parseHybridRejectsInvalidWeights() { + Map body = Map.of( + "query", "metrics", + "match_weight", 0.9, + "semantic_weight", 0.9); + assertThrows(SearchException.class, + () -> SearchArgs.fromBody(body, SearchQuery.Mode.HYBRID)); + } + + @Test + public void parseSemanticString() throws Exception { + SearchArgs.Semantic semantic = + (SearchArgs.Semantic) SearchArgs.fromBody("sales revenue", SearchQuery.Mode.SEMANTIC); + assertEquals("sales revenue", semantic.queryText()); + assertNull(semantic.field()); + } + + @Test + public void parseSemanticFieldAndQueryObject() throws Exception { + Map body = Map.of( + "field", MetastoreTableMapper.FIELD_SEARCH_TEXT, + "query", "orders"); + SearchArgs.Semantic semantic = + (SearchArgs.Semantic) SearchArgs.fromBody(body, SearchQuery.Mode.SEMANTIC); + assertEquals(MetastoreTableMapper.FIELD_SEARCH_TEXT, semantic.field()); + assertEquals("orders", semantic.queryText()); + } + + @Test + public void parseSemanticSingleFieldMap() throws Exception { + Map body = Map.of(MetastoreTableMapper.FIELD_SEARCH_TEXT, "inventory"); + SearchArgs.Semantic semantic = + (SearchArgs.Semantic) SearchArgs.fromBody(body, SearchQuery.Mode.SEMANTIC); + assertEquals(MetastoreTableMapper.FIELD_SEARCH_TEXT, semantic.field()); + assertEquals("inventory", semantic.queryText()); + } + + @Test + public void resolveHybridUsesSoleHybridField() throws Exception { + SearchArgs.Hybrid args = + (SearchArgs.Hybrid) SearchArgs.fromBody("sales revenue", SearchQuery.Mode.HYBRID); + HybridSearch.ResolvedHybridQuery resolved = HybridSearch.resolve(args, hybridMapping()); + assertEquals(MetastoreTableMapper.FIELD_SEARCH_TEXT, resolved.field()); + assertEquals("sales revenue", resolved.queryText()); + } + + @Test + public void resolveHybridRejectsNonHybridField() { + SearchArgs.Hybrid args = new SearchArgs.Hybrid( + "t1", MetastoreTableMapper.FIELD_TABLE, null, null); + assertThrows(SearchException.class, () -> HybridSearch.resolve(args, hybridMapping())); + } + + @Test + public void resolveHybridRequiresFieldWhenMultipleHybridFieldsExist() throws Exception { + Configuration conf = new Configuration(false); + Map fields = new LinkedHashMap<>(); + fields.put( + "field_a", + new FieldSchema.TextFieldSchema( + "field_a", new SearchParams(true, "model-a", SearchParams.VectorDistance.COSINE))); + fields.put( + "field_b", + new FieldSchema.TextFieldSchema( + "field_b", new SearchParams(true, "model-b", SearchParams.VectorDistance.COSINE))); + IndexMapping mapping = new IndexMapping("idx", conf, fields); + SearchArgs.Hybrid args = + (SearchArgs.Hybrid) SearchArgs.fromBody("query", SearchQuery.Mode.HYBRID); + + assertThrows(SearchException.class, () -> HybridSearch.resolve(args, mapping)); + } + + @Test + public void resolveSemanticUsesSoleSemanticField() throws Exception { + SearchArgs.Semantic args = + (SearchArgs.Semantic) SearchArgs.fromBody("sales revenue", SearchQuery.Mode.SEMANTIC); + SemanticSearch.ResolvedSemanticQuery resolved = + SemanticSearch.resolve(args, hybridMapping()); + assertEquals(MetastoreTableMapper.FIELD_SEARCH_TEXT, resolved.field()); + assertEquals("sales revenue", resolved.queryText()); + } + + @Test + public void resolveSemanticRejectsNonSemanticField() { + SearchArgs.Semantic args = new SearchArgs.Semantic("t1", MetastoreTableMapper.FIELD_TABLE); + assertThrows(SearchException.class, () -> SemanticSearch.resolve(args, hybridMapping())); + } + + @Test + public void serializeMatchRoundTrip() throws Exception { + assertRoundTrip(new SearchArgs.Match("orders"), SearchQuery.Mode.MATCH); + assertEquals(Map.of("query", "orders"), SearchArgs.toBody(new SearchArgs.Match("orders"))); + } + + @Test + public void serializeSemanticRoundTrip() throws Exception { + assertRoundTrip(new SearchArgs.Semantic("sales", null), SearchQuery.Mode.SEMANTIC); + assertRoundTrip( + new SearchArgs.Semantic("orders", MetastoreTableMapper.FIELD_SEARCH_TEXT), + SearchQuery.Mode.SEMANTIC); + assertEquals( + Map.of( + "query", "orders", + "field", MetastoreTableMapper.FIELD_SEARCH_TEXT), + SearchArgs.toBody( + new SearchArgs.Semantic("orders", MetastoreTableMapper.FIELD_SEARCH_TEXT))); + } + + @Test + public void serializeHybridRoundTrip() throws Exception { + assertRoundTrip(new SearchArgs.Hybrid("sales", null, null, null), SearchQuery.Mode.HYBRID); + assertRoundTrip( + new SearchArgs.Hybrid("metrics", MetastoreTableMapper.FIELD_SEARCH_TEXT, null, null), + SearchQuery.Mode.HYBRID); + assertRoundTrip( + new SearchArgs.Hybrid("metrics", null, 0.7f, 0.3f), + SearchQuery.Mode.HYBRID); + assertRoundTrip( + new SearchArgs.Hybrid("metrics", MetastoreTableMapper.FIELD_SEARCH_TEXT, 0.7f, 0.3f), + SearchQuery.Mode.HYBRID); + } + + @Test + public void serializeToQueryBodyRoundTrip() throws Exception { + SearchQuery query = SearchQuery.of("sales", SearchQuery.Mode.HYBRID, 10); + SearchQuery roundTripped = SearchQuery.fromQueryBody( + query.toQueryBody(), + query.mode(), + query.catalogName(), + query.databaseName(), + query.limit(), + query.returnFields()); + assertEquals(query.args(), roundTripped.args()); + assertEquals(query.mode(), roundTripped.mode()); + } + + private static void assertRoundTrip(SearchArgs args, SearchQuery.Mode mode) throws SearchException { + SearchArgs roundTripped = SearchArgs.fromBody(SearchArgs.toBody(args), mode); + assertEquals(args, roundTripped); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/search/TestSearchQuery.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/search/TestSearchQuery.java new file mode 100644 index 000000000000..80f341c504ed --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/search/TestSearchQuery.java @@ -0,0 +1,69 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.search; + +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hive.search.exception.SearchException; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +@Category(MetastoreUnitTest.class) +public class TestSearchQuery { + + @Test + public void ofTextDefaultsToHybridMode() throws Exception { + SearchQuery query = SearchQuery.of("sales"); + assertTrue(query.args() instanceof SearchArgs.Hybrid); + assertEquals("sales", ((SearchArgs.Hybrid) query.args()).queryText()); + assertEquals(SearchQuery.Mode.HYBRID, query.mode()); + assertEquals(0, query.limit()); + } + + @Test + public void ofTextWithModeAndLimit() throws Exception { + SearchQuery query = SearchQuery.of("sales", SearchQuery.Mode.SEMANTIC, 25); + assertTrue(query.args() instanceof SearchArgs.Semantic); + assertEquals(SearchQuery.Mode.SEMANTIC, query.mode()); + assertEquals(25, query.limit()); + } + + @Test + public void ofTableNameUsesKeywordMode() throws Exception { + SearchQuery query = SearchQuery.of("orders", "hive", "default"); + assertTrue(query.args() instanceof SearchArgs.Match); + assertEquals("orders", ((SearchArgs.Match) query.args()).queryText()); + assertEquals(SearchQuery.Mode.MATCH, query.mode()); + assertEquals("hive", query.catalogName()); + assertEquals("default", query.databaseName()); + } + + @Test + public void rejectsEmptyQueryText() { + assertThrows(SearchException.class, () -> SearchQuery.of("")); + } + + @Test + public void rejectsNegativeLimit() { + assertThrows(SearchException.class, + () -> SearchQuery.of("sales", SearchQuery.Mode.HYBRID, -1)); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/search/TestSemanticSearch.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/search/TestSemanticSearch.java new file mode 100644 index 000000000000..a2cde281ef4b --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/search/TestSemanticSearch.java @@ -0,0 +1,90 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.search; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest; +import org.apache.hive.search.exception.SearchException; +import org.apache.hive.search.mapping.FieldSchema; +import org.apache.hive.search.mapping.IndexMapping; +import org.apache.hive.search.mapping.SearchParams; +import org.apache.hive.search.metastore.MetastoreSchemas; +import org.apache.hive.search.metastore.MetastoreTableMapper; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.util.LinkedHashMap; +import java.util.Map; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; + +@Category(MetastoreUnitTest.class) +public class TestSemanticSearch { + + private static IndexMapping defaultMapping() { + Configuration conf = new Configuration(false); + return MetastoreSchemas.defaultHiveTablesMapping("hive_tables", "bge-small", conf); + } + + @Test + public void resolveStringUsesSoleSemanticField() throws Exception { + SemanticSearch.ResolvedSemanticQuery resolved = + SemanticSearch.resolve(new SearchArgs.Semantic("sales revenue", null), defaultMapping()); + assertEquals(MetastoreTableMapper.FIELD_SEARCH_TEXT, resolved.field()); + assertEquals("sales revenue", resolved.queryText()); + } + + @Test + public void resolveExplicitField() throws Exception { + SemanticSearch.ResolvedSemanticQuery resolved = SemanticSearch.resolve( + new SearchArgs.Semantic("orders", MetastoreTableMapper.FIELD_SEARCH_TEXT), + defaultMapping()); + assertEquals(MetastoreTableMapper.FIELD_SEARCH_TEXT, resolved.field()); + assertEquals("orders", resolved.queryText()); + } + + @Test + public void resolveRejectsNonSemanticField() { + assertThrows( + SearchException.class, + () -> + SemanticSearch.resolve( + new SearchArgs.Semantic("t1", MetastoreTableMapper.FIELD_TABLE), + defaultMapping())); + } + + @Test + public void resolveRequiresFieldWhenMultipleSemanticFieldsExist() { + Configuration conf = new Configuration(false); + Map fields = new LinkedHashMap<>(); + fields.put( + "field_a", + new FieldSchema.TextFieldSchema( + "field_a", new SearchParams(true, "model-a", SearchParams.VectorDistance.COSINE))); + fields.put( + "field_b", + new FieldSchema.TextFieldSchema( + "field_b", new SearchParams(true, "model-b", SearchParams.VectorDistance.COSINE))); + IndexMapping mapping = new IndexMapping("idx", conf, fields); + + assertThrows( + SearchException.class, + () -> SemanticSearch.resolve(new SearchArgs.Semantic("query", null), mapping)); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/testutil/InMemoryIndexStateClient.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/testutil/InMemoryIndexStateClient.java new file mode 100644 index 000000000000..72c5a3d8c8a6 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/testutil/InMemoryIndexStateClient.java @@ -0,0 +1,85 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.testutil; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.util.HashMap; +import java.util.Map; +import java.util.Optional; + +import org.apache.hive.search.exception.IndexIOException; +import org.apache.hive.search.index.manifest.IndexManifest; +import org.apache.hive.search.index.store.IndexStateClient; + +/** In-memory {@link IndexStateClient} for unit tests. */ +public final class InMemoryIndexStateClient implements IndexStateClient { + private final Map files = new HashMap<>(); + private byte[] stagingManifest; + + @Override + public Optional readManifest() throws IOException { + byte[] bytes = files.get(IndexManifest.MANIFEST_FILE_NAME); + return bytes == null ? Optional.empty() : Optional.of(IndexManifest.fromJson(bytes)); + } + + @Override + public InputStream read(String fileName) throws IOException { + byte[] bytes = files.get(fileName); + if (bytes == null) { + throw new IndexIOException("missing file: " + fileName); + } + return new ByteArrayInputStream(bytes); + } + + @Override + public void write(String fileName, InputStream stream) throws IOException { + files.put(fileName, stream.readAllBytes()); + } + + @Override + public void delete(String fileName) throws IOException { + files.remove(fileName); + } + + @Override + public Optional readStagingManifest() throws IOException { + return stagingManifest == null ? Optional.empty() : Optional.of(IndexManifest.fromJson(stagingManifest)); + } + + @Override + public void writeStagingManifest(IndexManifest manifest) throws IOException { + stagingManifest = manifest.toJsonBytes(); + } + + @Override + public void clearStagingManifest() { + stagingManifest = null; + } + + @Override + public void clear() { + files.clear(); + stagingManifest = null; + } + + public boolean hasFile(String fileName) { + return files.containsKey(fileName); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/testutil/IndexMutationApplier.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/testutil/IndexMutationApplier.java new file mode 100644 index 000000000000..c4dca5209737 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/testutil/IndexMutationApplier.java @@ -0,0 +1,82 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.testutil; + +import org.apache.hadoop.hive.common.DatabaseName; +import org.apache.hadoop.hive.common.TableName; +import org.apache.hadoop.hive.metastore.api.Table; +import org.apache.hive.search.exception.IndexException; +import org.apache.hive.search.index.Indexer; +import org.apache.hive.search.index.IndexManager; +import org.apache.hive.search.mapping.TableDocument; +import org.apache.hive.search.metastore.MetastoreEventListener; +import org.apache.hive.search.metastore.MetastoreTableMapper; + +import java.io.IOException; +import java.util.List; + +/** Applies {@link MetastoreEventListener.IndexTask} mutations to an {@link Indexer}. */ +public final class IndexMutationApplier { + private final IndexManager indexManager; + private final Indexer indexer; + + public IndexMutationApplier(IndexManager indexManager, Indexer indexer) { + this.indexManager = indexManager; + this.indexer = indexer; + } + + public void apply(MetastoreEventListener.IndexTask task) + throws IndexException, IOException { + if (!task.databasesToDrop.isEmpty()) { + indexer.deleteDatabases(task.databasesToDrop.toArray(new DatabaseName[0])); + } + if (!task.tablesToDrop.isEmpty()) { + String[] docIds = task.tablesToDrop.stream() + .map(MetastoreTableMapper::tableId) + .toList() + .toArray(new String[0]); + indexer.delete(docIds); + } + if (!task.tablesToAdd.isEmpty()) { + List newDocs = task.tablesToAdd.values().stream() + .map(table -> MetastoreTableMapper.fromTable(table, indexManager.mapping())) + .toList(); + indexer.addDocuments(newDocs); + } + } + + public void addTable(Table table) throws IndexException, IOException { + MetastoreEventListener.IndexTask task = new MetastoreEventListener.IndexTask(); + TableName tableName = new TableName(table.getCatName(), table.getDbName(), table.getTableName()); + task.tablesToAdd.put(tableName, table); + apply(task); + } + + public void dropTable(TableName tableName) throws IndexException, IOException { + MetastoreEventListener.IndexTask task = new MetastoreEventListener.IndexTask(); + task.tablesToDrop.add(tableName); + apply(task); + } + + public void replaceTable(Table before, Table after) throws IndexException, IOException { + MetastoreEventListener.IndexTask task = new MetastoreEventListener.IndexTask(); + task.tablesToDrop.add(new TableName(before.getCatName(), before.getDbName(), before.getTableName())); + task.tablesToAdd.put(new TableName(after.getCatName(), after.getDbName(), after.getTableName()), after); + apply(task); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/testutil/MetastoreBootstrapMocks.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/testutil/MetastoreBootstrapMocks.java new file mode 100644 index 000000000000..b3ef31b0689c --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/testutil/MetastoreBootstrapMocks.java @@ -0,0 +1,85 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.testutil; + +import org.apache.hadoop.hive.metastore.IMetaStoreClient; +import org.apache.hadoop.hive.metastore.api.CurrentNotificationEventId; +import org.apache.hadoop.hive.metastore.api.NotificationEventRequest; +import org.apache.hadoop.hive.metastore.api.NotificationEventResponse; +import org.apache.hadoop.hive.metastore.api.Table; + +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyList; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.when; + +/** Mockito helpers for bootstrap and indexer initialization tests. */ +public final class MetastoreBootstrapMocks { + private MetastoreBootstrapMocks() {} + + public static void stubCatchUp(IMetaStoreClient client) throws Exception { + NotificationEventResponse empty = new NotificationEventResponse(); + empty.setEvents(new ArrayList<>()); + when(client.getNextNotification(any(NotificationEventRequest.class), eq(false), eq(null))) + .thenReturn(empty); + } + + public static void stubCurrentNotificationId(IMetaStoreClient client, long eventId) + throws Exception { + when(client.getCurrentNotificationEventId()) + .thenReturn(new CurrentNotificationEventId(eventId)); + } + + public static void stubBootstrapCatalog(IMetaStoreClient client, Table... tables) + throws Exception { + Map> tablesByDb = new LinkedHashMap<>(); + for (Table table : tables) { + tablesByDb.computeIfAbsent(table.getDbName(), db -> new ArrayList<>()).add(table); + } + when(client.getAllDatabases()).thenReturn(new ArrayList<>(tablesByDb.keySet())); + for (Map.Entry> entry : tablesByDb.entrySet()) { + String db = entry.getKey(); + List tableNames = entry.getValue().stream() + .map(Table::getTableName) + .collect(Collectors.toList()); + when(client.getAllTables(db)).thenReturn(tableNames); + when(client.getTableObjectsByName(eq(db), anyList())).thenAnswer(invocation -> { + @SuppressWarnings("unchecked") + List requested = invocation.getArgument(1); + return requested.stream() + .map(name -> findTable(entry.getValue(), name)) + .collect(Collectors.toList()); + }); + } + } + + private static Table findTable(List

tables, String name) { + for (Table table : tables) { + if (table.getTableName().equals(name)) { + return table; + } + } + throw new IllegalArgumentException("Unknown table: " + name); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/testutil/RealMetastoreServer.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/testutil/RealMetastoreServer.java new file mode 100644 index 000000000000..c483104a8380 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/testutil/RealMetastoreServer.java @@ -0,0 +1,105 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.testutil; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.metastore.HiveMetaStoreClient; +import org.apache.hadoop.hive.metastore.IMetaStoreClient; +import org.apache.hadoop.hive.metastore.MetaStoreTestUtils; +import org.apache.hadoop.hive.metastore.client.builder.DatabaseBuilder; +import org.apache.hadoop.hive.metastore.client.builder.TableBuilder; +import org.apache.hadoop.hive.metastore.conf.MetastoreConf; +import org.apache.hadoop.hive.metastore.conf.MetastoreConf.ConfVars; +import org.apache.hadoop.hive.metastore.messaging.json.JSONMessageEncoder; +import org.apache.hadoop.hive.metastore.security.HadoopThriftAuthBridge; +import org.apache.hive.hcatalog.listener.DbNotificationListener; + +/** Starts an in-process HMS backed by Derby for integration tests. */ +public final class RealMetastoreServer implements AutoCloseable { + private static RealMetastoreServer instance; + + private final Configuration conf; + private final int port; + private final HiveMetaStoreClient client; + + private RealMetastoreServer(Configuration conf, int port, HiveMetaStoreClient client) { + this.conf = conf; + this.port = port; + this.client = client; + } + + public static synchronized RealMetastoreServer start() throws Exception { + if (instance != null) { + return instance; + } + Configuration conf = MetastoreConf.newMetastoreConf(); + MetaStoreTestUtils.setConfForStandloneMode(conf); + MetastoreConf.setBoolVar(conf, ConfVars.HIVE_IN_TEST, true); + MetastoreConf.setBoolVar(conf, ConfVars.CAPABILITY_CHECK, false); + MetastoreConf.setVar(conf, ConfVars.TRANSACTIONAL_EVENT_LISTENERS, + DbNotificationListener.class.getName()); + MetastoreConf.setVar(conf, ConfVars.EVENT_MESSAGE_FACTORY, JSONMessageEncoder.class.getName()); + MetastoreConf.setBoolVar(conf, ConfVars.HIVE_SUPPORT_CONCURRENCY, false); + int port = MetaStoreTestUtils.startMetaStoreWithRetry(HadoopThriftAuthBridge.getBridge(), conf); + HiveMetaStoreClient client = new HiveMetaStoreClient(conf); + instance = new RealMetastoreServer(conf, port, client); + return instance; + } + + public static synchronized RealMetastoreServer get() { + if (instance == null) { + throw new IllegalStateException("Call RealMetastoreServer.start() first"); + } + return instance; + } + + public Configuration conf() { + return new Configuration(conf); + } + + public IMetaStoreClient client() { + return client; + } + + public void createDatabase(String dbName) throws Exception { + new DatabaseBuilder().setName(dbName).create(client, conf); + } + + public void createTable(String dbName, String tableName, String comment) throws Exception { + new TableBuilder() + .setDbName(dbName) + .setTableName(tableName) + .addCol("id", "int") + .addTableParam("comment", comment) + .create(client, conf); + } + + public void dropTable(String dbName, String tableName) throws Exception { + client.dropTable(dbName, tableName); + } + + @Override + public synchronized void close() { + if (instance == null) { + return; + } + client.close(); + MetaStoreTestUtils.close(port); + instance = null; + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/testutil/StubEmbedModel.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/testutil/StubEmbedModel.java new file mode 100644 index 000000000000..e775fb555dee --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/testutil/StubEmbedModel.java @@ -0,0 +1,85 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.testutil; + +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.concurrent.atomic.AtomicInteger; + +import org.apache.hive.search.exception.IndexException; +import org.apache.hive.search.inference.EmbedModel; + +/** Deterministic embedding model for tests (no ONNX). */ +public final class StubEmbedModel implements EmbedModel { + private static final int DIMENSION = 8; + + private final String name; + private final AtomicInteger encodeBatchCalls = new AtomicInteger(); + + public StubEmbedModel(String name) { + this.name = name; + } + + public int encodeBatchCalls() { + return encodeBatchCalls.get(); + } + + @Override + public String name() { + return name; + } + + @Override + public float[] embed(TaskType task, String text) { + return vector(text + ":" + task.name()); + } + + @Override + public void close() { + // no-op + } + + private static float[] vector(String text) { + byte[] bytes = text.getBytes(StandardCharsets.UTF_8); + float[] embedding = new float[DIMENSION]; + for (int i = 0; i < DIMENSION; i++) { + embedding[i] = (bytes[i % bytes.length] & 0xFF) / 255.0f; + } + float norm = 0f; + for (float value : embedding) { + norm += value * value; + } + norm = (float) Math.sqrt(norm); + if (norm > 0f) { + for (int i = 0; i < embedding.length; i++) { + embedding[i] /= norm; + } + } + return embedding; + } + + public static float[] queryVector(String text) { + return vector(text + ":" + TaskType.QUERY.name()); + } + + @Override + public float[][] embedBatch(TaskType task, String[] texts) throws IndexException { + encodeBatchCalls.incrementAndGet(); + return Arrays.stream(texts).map(text -> embed(task, text)).toArray(float[][]::new); + } +} diff --git a/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/testutil/TestLeaderElection.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/testutil/TestLeaderElection.java new file mode 100644 index 000000000000..e9a298d1d598 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/testutil/TestLeaderElection.java @@ -0,0 +1,73 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hive.search.testutil; + +import java.util.ArrayList; +import java.util.List; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.hive.metastore.leader.LeaderElection; +import org.apache.hadoop.hive.metastore.leader.LeaderException; + +/** Fixed leader election for metastore-search unit and integration tests. */ +public final class TestLeaderElection implements LeaderElection { + private final boolean leader; + private final List listeners = new ArrayList<>(); + private String name; + + public TestLeaderElection(boolean leader) { + this.leader = leader; + } + + @Override + public void tryBeLeader(Configuration conf, Object mutex) throws LeaderException { + if (leader) { + for (LeadershipStateListener listener : listeners) { + try { + listener.takeLeadership(this); + } catch (Exception e) { + throw new LeaderException("failed to notify test leader listener", e); + } + } + } + } + + @Override + public boolean isLeader() { + return leader; + } + + @Override + public void addStateListener(LeadershipStateListener listener) { + listeners.add(listener); + } + + @Override + public void setName(String name) { + this.name = name; + } + + @Override + public String getName() { + return name; + } + + @Override + public void close() { + } +} diff --git a/standalone-metastore/packaging/pom.xml b/standalone-metastore/packaging/pom.xml index 4fb9de38f3ea..6b8f7a311e5a 100644 --- a/standalone-metastore/packaging/pom.xml +++ b/standalone-metastore/packaging/pom.xml @@ -189,6 +189,11 @@ + + org.apache.hive + hive-standalone-metastore-search + ${project.version} + org.apache.hive hive-exec diff --git a/standalone-metastore/packaging/src/assembly/src.xml b/standalone-metastore/packaging/src/assembly/src.xml index 2b202bd15097..ce1ed4b96f9c 100644 --- a/standalone-metastore/packaging/src/assembly/src.xml +++ b/standalone-metastore/packaging/src/assembly/src.xml @@ -43,6 +43,7 @@ metastore-server/**/* metastore-tools/**/* metastore-rest-catalog/**/* + metastore-search/**/* packaging/**/* checkstyle/**/* spotbugs/**/* diff --git a/standalone-metastore/pom.xml b/standalone-metastore/pom.xml index dc25fe3410a8..07e467fd7de4 100644 --- a/standalone-metastore/pom.xml +++ b/standalone-metastore/pom.xml @@ -31,6 +31,7 @@ metastore-server metastore-tools metastore-rest-catalog + metastore-search packaging