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/standalone-metastore/metastore-search/pom.xml b/standalone-metastore/metastore-search/pom.xml new file mode 100644 index 000000000000..948357ef80ab --- /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.16.2 + 1.16.1-beta26 + 10.4.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 + + + dev.langchain4j + langchain4j + ${langchain4j.version} + + + dev.langchain4j + langchain4j-embeddings-bge-small-en-v15 + ${langchain4j.embeddings.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/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/IndexStateConfig.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/config/IndexStateConfig.java new file mode 100644 index 000000000000..f30a82131cb1 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/config/IndexStateConfig.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 IndexStateConfig(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..89104bc2e69a --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/config/InferenceConfig.java @@ -0,0 +1,130 @@ +/* + * 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 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 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..4def067e0a93 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/config/SearchConfig.java @@ -0,0 +1,86 @@ +/* + * 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; + +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() { + float weight = configuration.getFloat(HYBRID_SEMANTIC_WEIGHT, HYBRID_SEMANTIC_WEIGHT_DEFAULT); + if (weight >= 1 || weight <= 0) { + throw new IllegalArgumentException("metastore.search.hybrid.semantic.weight " + + "should be configured as between 0.0 and 1.0"); + } + return weight; + } + + 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..3fb795b840ab --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/IndexManager.java @@ -0,0 +1,184 @@ +/* + * 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.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.IndexStateConfig; +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; + + 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 { + IndexStateConfig 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]; + } + } + + public void checkIndexState() throws IndexNotHealthyException { + if (exception != null) { + throw exception; + } + } + + public IndexMapping mapping() { + return mapping; + } + + public Directory directory() { + return directory; + } + + @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..9082112f462e --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/Indexer.java @@ -0,0 +1,306 @@ +/* + * 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 SnapshotDeletionPolicy snapshotter; + private FlushTrackingWriter writer; + private int commitFlushThreshold; + + public Indexer(IndexManager index, EmbedModelRegistry registry) { + this.indexManager = index; + this.modelRegistry = registry; + this.embeddingCache = registry.embeddingCache(); + this.commitFlushThreshold = + new IndexConfig(index.mapping().configuration()).getCommitFlushes(); + } + + /** Package-private for bootstrap to apply a different commit threshold temporarily. */ + public void setCommitFlushThreshold(int commitFlushThreshold) { + this.commitFlushThreshold = Math.max(0, commitFlushThreshold); + } + + 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.encodeBatch(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..c70ecc5d353c --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/index/manifest/IndexManifest.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.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..697e0f22f2ca --- /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.IndexStateConfig; +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 { + IndexStateConfig.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..5a7bf3ac9f2d --- /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[] encode(TaskType task, String text) throws IndexException; + + default float[][] encodeBatch(TaskType task, String[] texts) throws IndexException { + float[][] result = new float[texts.length][]; + for (int i = 0; i < texts.length; i++) { + result[i] = encode(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..ea61af45890c --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/EmbedModelRegistry.java @@ -0,0 +1,78 @@ +/* + * 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(); + InferenceConfig.EmbeddingModelSpec embeddingConfig = inference.embedding(); + String modelName = inference.modelName(); + EmbedModel embedModel = new LocalOnnxEmbeddingModel(modelName, embeddingConfig.getModelDir(), + embeddingConfig.getPrompt()); + long warmupStart = System.currentTimeMillis(); + try { + embedModel.encode(EmbedModel.TaskType.QUERY, "warmup"); + } catch (IndexException e) { + throw new InitializeException("Failed to warm up embedding model '" + modelName + "'", e); + } + LOG.info("Loaded embedding model '{}' from {} in {}ms (warmup {}ms)", + modelName, embeddingConfig.getModelDir(), 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/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/LocalOnnxEmbeddingModel.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/LocalOnnxEmbeddingModel.java new file mode 100644 index 000000000000..e115e18132ba --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/inference/LocalOnnxEmbeddingModel.java @@ -0,0 +1,106 @@ +/* + * 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 dev.langchain4j.data.embedding.Embedding; +import dev.langchain4j.data.segment.TextSegment; +import dev.langchain4j.model.embedding.onnx.OnnxEmbeddingModel; +import dev.langchain4j.model.embedding.onnx.PoolingMode; + +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; + +import org.apache.hive.search.config.InferenceConfig; +import org.apache.hive.search.exception.IndexException; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public final class LocalOnnxEmbeddingModel implements EmbedModel { + private static final Logger LOG = LoggerFactory.getLogger(LocalOnnxEmbeddingModel.class); + + private final String name; + private final OnnxEmbeddingModel model; + private final EmbeddingPrompt prompt; + private final ExecutorService inferExecutor; + + public LocalOnnxEmbeddingModel(String name, Path modelDir, EmbeddingPrompt prompt) { + this.name = name; + this.prompt = prompt == null ? EmbeddingPrompt.none() : prompt; + this.inferExecutor = Executors.newSingleThreadExecutor(r -> { + Thread thread = new Thread(r, "EmbedModel-" + name); + thread.setDaemon(true); + return thread; + }); + this.model = new OnnxEmbeddingModel( + modelDir.resolve(InferenceConfig.MODEL_ONNX_FILE), + modelDir.resolve(InferenceConfig.TOKENIZER), + PoolingMode.MEAN, + inferExecutor); + } + + @Override + public String name() { + return name; + } + + @Override + public float[] encode(TaskType task, String text) throws IndexException { + try { + return model.embed(prompt.prefixFor(task) + text).content().vector(); + } catch (RuntimeException e) { + throw IndexException.wrap("Failed to encode text with model '" + name + "'", e); + } + } + + @Override + public float[][] encodeBatch(TaskType task, String[] texts) throws IndexException { + try { + String prefix = prompt.prefixFor(task); + List segments = new ArrayList<>(texts.length); + for (String text : texts) { + segments.add(TextSegment.from(prefix + text)); + } + List embeddings = model.embedAll(segments).content(); + float[][] vectors = new float[texts.length][]; + for (int i = 0; i < texts.length; i++) { + vectors[i] = embeddings.get(i).vector(); + } + return vectors; + } catch (RuntimeException e) { + throw IndexException.wrap("Failed to encode batch with model '" + name + "'", e); + } + } + + @Override + public void close() { + inferExecutor.shutdown(); + try { + if (!inferExecutor.awaitTermination(30, TimeUnit.SECONDS)) { + inferExecutor.shutdownNow(); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + inferExecutor.shutdownNow(); + } + 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..97f803dc9fbe --- /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.IndexStateConfig; +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 IndexStateConfig store() { + return new IndexStateConfig(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..df997b485639 --- /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 {@lin}. */ +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..5aad15051fc0 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/MetastoreEventHandler.java @@ -0,0 +1,382 @@ +/* + * 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(); + 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..3cb7d317a1d3 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/MetastoreIndexer.java @@ -0,0 +1,355 @@ +/* + * 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(); + } + + 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..917755572a97 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/metastore/MetastoreTableMapper.java @@ -0,0 +1,197 @@ +/* + * 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 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)); + /** Max commented data columns included in search-oriented column fields. */ + static final int MAX_SEARCH_COLUMNS = 100; + + 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 = + table.getSd() == null ? "" : nullToEmpty(table.getSd().getLocation()); + String comment = ""; + if (table.getParameters() != null && table.getParameters().get("comment") != null) { + comment = nullToEmpty(table.getParameters().get("comment")); + } + String columns = formatColumnsForStorage(table); + String columnNames = formatColumnNamesForSearch(table); + String columnComments = formatColumnCommentsForSearch(table); + String searchText = buildSearchText(name, comment, formatColumnsForSearch(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); + } + + private static String buildSearchText(String tableName, String comment, String searchColumns) { + String normalizedTableName = tableName.toLowerCase(Locale.ROOT); + if (comment.isEmpty()) { + if (searchColumns.isEmpty()) { + return normalizedTableName; + } + return normalizedTableName + " " + searchColumns; + } + if (searchColumns.isEmpty()) { + return normalizedTableName + " " + comment; + } + return normalizedTableName + " " + comment + " " + searchColumns; + } + + 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 commented = new ArrayList<>(); + for (FieldSchema column : table.getSd().getCols()) { + if (StringUtils.isNotEmpty(column.getComment())) { + commented.add(column); + } + } + int limit = Math.min(commented.size(), MAX_SEARCH_COLUMNS); + List parts = new ArrayList<>(limit); + for (int i = 0; i < limit; i++) { + parts.add(commented.get(i).getComment()); + } + String formatted = String.join("; ", parts); + if (commented.size() > MAX_SEARCH_COLUMNS) { + return formatted + "; ... (+" + (commented.size() - MAX_SEARCH_COLUMNS) + " more)"; + } + return formatted; + } + + private static String formatColumnsForSearch(Table table) { + if (table.getSd() == null || table.getSd().getCols() == null) { + return ""; + } + List commented = new ArrayList<>(); + for (FieldSchema column : table.getSd().getCols()) { + if (StringUtils.isNotEmpty(column.getComment())) { + commented.add(column); + } + } + int limit = Math.min(commented.size(), MAX_SEARCH_COLUMNS); + List parts = new ArrayList<>(limit); + for (int i = 0; i < limit; i++) { + parts.add(formatColumnForSearch(commented.get(i))); + } + String formatted = String.join("; ", parts); + if (commented.size() > MAX_SEARCH_COLUMNS) { + return formatted + "; ... (+" + (commented.size() - MAX_SEARCH_COLUMNS) + " more)"; + } + return formatted; + } + + 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) { + return column.getName() + " " + column.getComment(); + } + + 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..74fb99d18b78 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/HybridSearch.java @@ -0,0 +1,187 @@ +/* + * 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.config.SearchConfig; +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.metastore.MetastoreTableMapper; + +/** Builds RRF queries that fuse lexical and semantic retrieval on hybrid fields. */ +public final class HybridSearch { + private HybridSearch() {} + + public record ParsedHybridQuery( + String field, + String queryText, + Float semanticWeight) {} + + public static ParsedHybridQuery parse(Object hybridBody, IndexMapping mapping) + throws SearchException { + if (hybridBody instanceof String queryText) { + String field = requireHybridField(mapping, null); + return new ParsedHybridQuery(field, queryText, null); + } + if (!(hybridBody instanceof Map body)) { + throw new SearchException("hybrid query must be a string or object"); + } + @SuppressWarnings("unchecked") + Map hybridMap = (Map) body; + + HybridWeights weights = parseHybridWeights(hybridMap); + + if (hybridMap.containsKey("field") && hybridMap.containsKey("query")) { + String field = hybridMap.get("field").toString(); + validateHybridField(mapping, field); + return new ParsedHybridQuery( + field, hybridMap.get("query").toString(), weights.getSemanticWeight()); + } + + if (hybridMap.containsKey("query")) { + String field = requireHybridField(mapping, null); + return new ParsedHybridQuery( + field, hybridMap.get("query").toString(), weights.getSemanticWeight()); + } + + List> fieldEntries = + hybridMap.entrySet().stream() + .filter(entry -> !isHybridOption(entry.getKey())) + .filter(entry -> entry.getValue() != null) + .toList(); + if (fieldEntries.size() == 1) { + Map.Entry entry = fieldEntries.getFirst(); + validateHybridField(mapping, entry.getKey()); + return new ParsedHybridQuery( + entry.getKey(), entry.getValue().toString(), weights.getSemanticWeight()); + } + + throw new SearchException( + "hybrid query must be a string, {field:, query:}, {query:}, or {: text}"); + } + + public static SearchInternal.FusionRequest toFusionRequest( + ParsedHybridQuery hybrid, int defaultSize, SearchConfig searchConfig) { + float semanticWeight = hybrid.semanticWeight() != null ? + hybrid.semanticWeight() : searchConfig.getHybridSemanticWeight(); + float matchWeight = 1.0f - semanticWeight; + Map lexicalQuery = lexicalQueryForHybrid(hybrid); + List retrievers = + List.of( + new SearchInternal.RetrieverSpec(lexicalQuery, matchWeight, "match"), + new SearchInternal.RetrieverSpec( + Map.of("semantic", Map.of(hybrid.field(), hybrid.queryText())), semanticWeight, "semantic")); + return new SearchInternal.FusionRequest(retrievers, defaultSize); + } + + private static Map lexicalQueryForHybrid(ParsedHybridQuery hybrid) { + if (MetastoreTableMapper.FIELD_SEARCH_TEXT.equals(hybrid.field())) { + return Map.of("table_keyword", hybrid.queryText()); + } + return Map.of("match", Map.of(hybrid.field(), hybrid.queryText())); + } + + private record HybridWeights(Float match, Float semantic) { + HybridWeights validate() { + if (!validate(match)) { + throw new IllegalArgumentException(); + } + if (!validate(semantic)) { + throw new IllegalArgumentException(); + } + if (match != null && semantic != null && (match + semantic) > 1.0f) { + throw new IllegalArgumentException(); + } + return this; + } + + boolean validate(Float weight) { + if (weight == null) { + return true; + } + if (weight >= 1f || weight <= 0.0f) { + return false; + } + return true; + } + + float getSemanticWeight() { + if (semantic != null) { + return semantic; + } else if (match != null) { + return 1.0f - match; + } + return 0.4f; + } + } + + private static HybridWeights parseHybridWeights(Map hybridMap) { + Float match = null; + Float semantic = null; + if (hybridMap.get("match_weight") instanceof Number matchWeight) { + match = matchWeight.floatValue(); + } + if (hybridMap.get("semantic_weight") instanceof Number semanticWeight) { + semantic = semanticWeight.floatValue(); + } + if (hybridMap.get("weights") instanceof Map weights) { + if (weights.get("match") instanceof Number matchValue) { + match = matchValue.floatValue(); + } + if (weights.get("semantic") instanceof Number semanticValue) { + semantic = semanticValue.floatValue(); + } + } + return new HybridWeights(match, semantic).validate(); + } + + private static boolean isHybridOption(String key) { + return "match_weight".equals(key) + || "semantic_weight".equals(key) + || "weights".equals(key); + } + + 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..e16173ab62b8 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/LuceneSearchBackend.java @@ -0,0 +1,151 @@ +/* + * 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.Map; +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.hadoop.hive.common.TableName; +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 = buildReturnFields(query); + int limit = query.limit() > 0 ? query.limit() : searchConfig.getDefaultLimit(); + try (SearchInternal searcher = session.getSearcher()) { + SearchReqResp.Response response = searcher.search(SearchReqResp.Request.validated( + toInternalQuery(query), + fields, + limit, + query.catalogName(), + query.databaseName())); + List hits = new ArrayList<>(); + for (Map raw : response.hits()) { + hits.add(toHit(raw, query.catalogName())); + } + return new TableSearchResult(hits, response.total()); + } + } + + @Override + public void close() throws Exception { + if (session != null) { + session.close(); + } + } + + private static List buildReturnFields(SearchQuery query) { + if (!query.returnFields().isEmpty()) { + return query.returnFields(); + } + return List.of( + MetastoreTableMapper.FIELD_DB, + MetastoreTableMapper.FIELD_TABLE, + MetastoreTableMapper.FIELD_OWNER, + MetastoreTableMapper.FIELD_COMMENT); + } + + private static Map toInternalQuery(SearchQuery query) { + String text = query.queryText(); + return switch (query.mode()) { + case KEYWORD -> Map.of("table_keyword", text); + case SEMANTIC -> Map.of("semantic", + Map.of(MetastoreTableMapper.FIELD_SEARCH_TEXT, text)); + case HYBRID -> Map.of("hybrid", text); + }; + } + + private static TableSearchHit toHit(Map raw, String defaultCatalog) { + TableName tableName = TableName.fromString(raw.get("_id").toString(), defaultCatalog, "default"); + float score = + raw.get("_score") instanceof Number number ? number.floatValue() : 0f; + Map fields = new LinkedHashMap<>(); + for (Map.Entry entry : raw.entrySet()) { + if ("_score".equals(entry.getKey()) || "_id".equals(entry.getKey())) { + continue; + } + if (entry.getValue() != null) { + fields.put(entry.getKey(), entry.getValue().toString()); + } + } + return new TableSearchHit(tableName, score, fields); + } +} 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..f2ad024272cf --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/SearchInternal.java @@ -0,0 +1,342 @@ +/* + * 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.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; + + public SearchInternal(SearcherManager manager, + IndexManager indexManager, + EmbedModelRegistry registry, + SearchConfig searchConfig, + BayesianScoreEstimator.Parameters parameters) throws IOException { + this.searcherManager = manager; + this.searcher = manager.acquire(); + this.mapping = indexManager.mapping(); + this.searchConfig = searchConfig; + this.parameters = parameters; + this.modelRegistry = Objects.requireNonNull(registry, "Model registry"); + } + + public record RetrieverSpec(Map query, float weight, String name) { + } + + public record FusionRequest(List retrievers, int size) {} + + public SearchReqResp.Response search(SearchReqResp.Request request) + throws SearchException, IOException { + TopDocs topDocs = executeQuery(request); + List> hits = new ArrayList<>(); + for (ScoreDoc scoreDoc : topDocs.scoreDocs) { + hits.add(readHit(scoreDoc, request.fields())); + } + return new SearchReqResp.Response(hits, topDocs.totalHits.value()); + } + + private TopDocs executeQuery(SearchReqResp.Request request) + throws SearchException, IOException { + int size = request.size() > 0 ? request.size() : searchConfig.getDefaultLimit(); + Query query = compileRootQuery(request.query(), size); + query = applyScopeFilter(query, request.catalogName(), request.databaseName()); + return searcher.search(query, size); + } + + private Query compileRootQuery(Map queryNode, int size) + throws SearchException, IOException { + if (queryNode.containsKey("hybrid")) { + return compileHybridQuery(queryNode.get("hybrid"), size); + } + return compileQuery(queryNode, size); + } + + private Query compileHybridQuery(Object hybridBody, int size) + throws SearchException, IOException { + HybridSearch.ParsedHybridQuery hybrid = HybridSearch.parse(hybridBody, mapping); + return compileFusionQuery(HybridSearch.toFusionRequest(hybrid, 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(Map queryNode, int knnK) + throws SearchException, IOException { + if (queryNode.containsKey("table_keyword")) { + return compileTableKeywordQuery(queryNode.get("table_keyword").toString()); + } + if (queryNode.containsKey("semantic")) { + return compileSemanticQuery(queryNode, knnK); + } + if (queryNode.containsKey("hybrid")) { + throw new SearchException("nested hybrid queries are not supported"); + } + throw new SearchException("supported queries: table_keyword, match, semantic, hybrid"); + } + + private Query compileTableKeywordQuery(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( + Map queryNode, int knnK) throws SearchException { + Object semanticBody = queryNode.get("semantic"); + String field; + String queryText; + if (semanticBody instanceof String text) { + field = requireSemanticField(null); + queryText = text; + } else if (semanticBody instanceof Map body) { + @SuppressWarnings("unchecked") + Map semanticMap = (Map) body; + if (semanticMap.containsKey("field") && semanticMap.containsKey("query")) { + field = semanticMap.get("field").toString(); + queryText = semanticMap.get("query").toString(); + } else if (semanticMap.size() == 1) { + Map.Entry entry = semanticMap.entrySet().iterator().next(); + field = entry.getKey(); + queryText = entry.getValue().toString(); + } else { + throw new SearchException("semantic query must be {field: text} or {field:, query:}"); + } + } else { + throw new SearchException("semantic query must be a string or object"); + } + + FieldSchema schema = mapping.fieldSchema(field); + if (!(schema instanceof FieldSchema.TextFieldSchema text) || !text.search().semantic()) { + throw new SearchException("field '" + field + "' is not semantically searchable"); + } + float[] embedding; + try { + embedding = modelRegistry.get(text.search().semanticModel()) + .encode(EmbedModel.TaskType.QUERY, queryText); + } catch (IndexException e) { + throw new SearchException( + "Failed to encode semantic query for field '" + field + "'", e); + } + if (embedding == null) { + throw new SearchException( + "embedding model '" + text.search().semanticModel() + "' returned null vector"); + } + return new KnnFloatVectorQuery(field, embedding, knnK); + } + + private String requireSemanticField(String field) throws SearchException { + if (StringUtils.isNotEmpty(field)) { + validateSemanticField(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 void validateSemanticField(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"); + } + } + + private Map readHit(ScoreDoc scoreDoc, List fields) + throws IOException { + Document stored = searcher.storedFields().document(scoreDoc.doc); + Map hit = new LinkedHashMap<>(); + hit.put("_score", scoreDoc.score); + 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) { + hit.put(field, values[0].stringValue()); + } + } + IndexableField[] idValues = stored.getFields("_id"); + if (idValues != null && idValues.length > 0) { + hit.put("_id", idValues[0].stringValue()); + } + return hit; + } + + @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..7d2432dec54b --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/SearchQuery.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.search; + +import java.util.List; + +import org.apache.commons.lang3.StringUtils; +import org.apache.hive.search.exception.SearchException; +import org.apache.hadoop.hive.common.TableName; + +/** Narrow public search contract (Thrift v1 semantics). */ +public record SearchQuery( + String queryText, + SearchReqResp.Mode mode, + String catalogName, + String databaseName, + int limit, + List returnFields) { + + public SearchQuery { + mode = mode == null ? SearchReqResp.Mode.HYBRID : mode; + returnFields = returnFields == null ? List.of() : List.copyOf(returnFields); + } + + 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 of(String queryText) throws SearchException { + validate(queryText, 0); + return new SearchQuery(queryText, SearchReqResp.Mode.HYBRID, null, null, 0, List.of()); + } + + public static SearchQuery of(String queryText, SearchReqResp.Mode mode, int limit) + throws SearchException { + validate(queryText, limit); + return new SearchQuery(queryText, mode, null, null, limit, List.of()); + } + + public static SearchQuery of(TableName tableName) throws SearchException { + validate(tableName.getTable(), 0); + return new SearchQuery(tableName.getTable(), SearchReqResp.Mode.KEYWORD, + tableName.getCat(), tableName.getDb(), 0, List.of()); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/SearchReqResp.java b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/SearchReqResp.java new file mode 100644 index 000000000000..a239879688b0 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/SearchReqResp.java @@ -0,0 +1,55 @@ +/* + * 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.hive.search.exception.SearchException; + +public final class SearchReqResp { + /** Public search mode exposed on the Metastore Thrift API. */ + public enum Mode { + KEYWORD, + SEMANTIC, + HYBRID + } + + public record Request( + Map query, + List fields, + int size, + String catalogName, + String databaseName) { + public Request { + query = query == null ? Map.of() : Map.copyOf(query); + fields = fields == null ? List.of() : List.copyOf(fields); + } + + public static Request validated(Map query, List fields, int size, + String catalogName, String databaseName) throws SearchException { + if (size < 0) { + throw new SearchException("size must be non-negative"); + } + return new Request(query, fields, size, catalogName, databaseName); + } + } + + public record Response(List> hits, long total) {} + +} 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..334db1ae7e8f --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/hive/search/search/TableSearchResult.java @@ -0,0 +1,27 @@ +/* + * 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) { + + public TableSearchResult { + hits = hits == null ? List.of() : List.copyOf(hits); + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/lucene/search/BayesianScoreEstimator.java b/standalone-metastore/metastore-search/src/main/java/org/apache/lucene/search/BayesianScoreEstimator.java new file mode 100644 index 000000000000..d13821214187 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/lucene/search/BayesianScoreEstimator.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.lucene.search; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Random; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.MultiTerms; +import org.apache.lucene.index.Term; +import org.apache.lucene.index.Terms; +import org.apache.lucene.index.TermsEnum; +import org.apache.lucene.util.BytesRef; + +/** + * Estimates {@link BayesianScoreQuery} parameters (alpha, beta, base rate) from corpus statistics + * via pseudo-query sampling. + * + *

The estimation algorithm: + * + *

    + *
  1. Reservoir-sample terms from the target field's indexed vocabulary + *
  2. Partition the sampled terms into pseudo-queries + *
  3. Run each pseudo-query via BM25 and collect the score distribution + *
  4. Estimate: beta = median(scores), alpha = 1 / std(scores) + *
  5. Estimate base rate: mean fraction of documents scoring above the 95th percentile + *
+ * + * @lucene.experimental + */ +public class BayesianScoreEstimator { + + /** Estimated parameters for {@link BayesianScoreQuery}. */ + public record Parameters(float alpha, float beta, float baseRate) {} + + private static final int DEFAULT_N_SAMPLES = 50; + private static final int DEFAULT_TOKENS_PER_QUERY = 5; + private static final double PERCENTILE_THRESHOLD = 0.95; + private static final float BASE_RATE_MIN = 1e-6f; + private static final float BASE_RATE_MAX = 0.5f; + + private BayesianScoreEstimator() {} + + /** + * Estimates BayesianScoreQuery parameters from the given index. + * + * @param searcher the index searcher to sample from + * @param field the indexed text field to create pseudo-queries for + * @param nSamples number of pseudo-queries to sample (default 50) + * @param tokensPerQuery number of indexed terms per pseudo-query (default 5) + * @param seed random seed for reproducible sampling + * @return estimated alpha, beta, and base rate + * @throws IOException if an I/O error occurs reading the index + */ + public static Parameters estimate( + IndexSearcher searcher, String field, int nSamples, int tokensPerQuery, long seed) + throws IOException { + if (nSamples <= 0) { + throw new IllegalArgumentException("nSamples must be positive, got " + nSamples); + } + if (tokensPerQuery <= 0) { + throw new IllegalArgumentException("tokensPerQuery must be positive, got " + tokensPerQuery); + } + + IndexReader reader = searcher.getIndexReader(); + int maxDoc = reader.maxDoc(); + if (maxDoc == 0) { + return new Parameters(1.0f, 0.0f, 0.01f); + } + + Random rng = new Random(seed); + List sampledTerms = + sampleVocabularyTerms(reader, field, Math.multiplyExact(nSamples, tokensPerQuery), rng); + if (sampledTerms.isEmpty()) { + return new Parameters(1.0f, 0.0f, 0.01f); + } + + // Create pseudo-queries from indexed vocabulary terms and collect scores. + List allScoreArrays = new ArrayList<>(); + List baseRateFractions = new ArrayList<>(); + + for (int offset = 0; offset < sampledTerms.size(); offset += tokensPerQuery) { + BooleanQuery.Builder builder = new BooleanQuery.Builder(); + int end = Math.min(offset + tokensPerQuery, sampledTerms.size()); + for (int i = offset; i < end; i++) { + builder.add( + new TermQuery(new Term(field, sampledTerms.get(i))), BooleanClause.Occur.SHOULD); + } + Query pseudoQuery = builder.build(); + + // Collect all scores + float[] scores = collectScores(searcher, pseudoQuery, maxDoc); + if (scores.length == 0) { + continue; + } + allScoreArrays.add(scores); + + // Base rate: fraction of docs above 95th percentile + float[] sorted = scores.clone(); + Arrays.sort(sorted); + int pIdx = (int) (sorted.length * PERCENTILE_THRESHOLD); + pIdx = Math.min(pIdx, sorted.length - 1); + float threshold = sorted[pIdx]; + int highCount = 0; + for (float s : scores) { + if (s >= threshold) { + highCount++; + } + } + baseRateFractions.add((float) highCount / maxDoc); + } + + if (allScoreArrays.isEmpty()) { + return new Parameters(1.0f, 0.0f, 0.01f); + } + + // Flatten all scores for global statistics + int totalScores = 0; + for (float[] arr : allScoreArrays) { + totalScores += arr.length; + } + float[] allScores = new float[totalScores]; + int offset = 0; + for (float[] arr : allScoreArrays) { + System.arraycopy(arr, 0, allScores, offset, arr.length); + offset += arr.length; + } + + // beta = median + Arrays.sort(allScores); + float beta = allScores[allScores.length / 2]; + + // alpha = 1 / std + double mean = 0; + for (float s : allScores) { + mean += s; + } + mean /= allScores.length; + double variance = 0; + for (float s : allScores) { + double diff = s - mean; + variance += diff * diff; + } + variance /= allScores.length; + double std = Math.sqrt(variance); + float alpha = std > 0 ? (float) (1.0 / std) : 1.0f; + + // base rate = mean of per-query fractions, clamped + float baseRate = 0; + for (float f : baseRateFractions) { + baseRate += f; + } + baseRate /= baseRateFractions.size(); + baseRate = Math.clamp(baseRate, BASE_RATE_MIN, BASE_RATE_MAX); + + return new Parameters(alpha, beta, baseRate); + } + + /** + * Estimates parameters with default settings (50 samples, 5 tokens per query, seed 42). + * + * @param searcher the index searcher + * @param field the text field + * @return estimated parameters + * @throws IOException if an I/O error occurs + */ + public static Parameters estimate(IndexSearcher searcher, String field) throws IOException { + return estimate(searcher, field, DEFAULT_N_SAMPLES, DEFAULT_TOKENS_PER_QUERY, 42); + } + + static List sampleVocabularyTerms( + IndexReader reader, String field, int sampleSize, Random rng) throws IOException { + Terms terms = MultiTerms.getTerms(reader, field); + if (terms == null) { + return new ArrayList<>(); + } + + List reservoir = new ArrayList<>(sampleSize); + TermsEnum termsEnum = terms.iterator(); + BytesRef term; + long seen = 0; + while ((term = termsEnum.next()) != null) { + seen++; + if (reservoir.size() < sampleSize) { + reservoir.add(BytesRef.deepCopyOf(term)); + } else { + long replacement = nextLong(rng, seen); + if (replacement < sampleSize) { + reservoir.set((int) replacement, BytesRef.deepCopyOf(term)); + } + } + } + return reservoir; + } + + private static long nextLong(Random rng, long bound) { + long bits; + long value; + do { + bits = rng.nextLong() >>> 1; + value = bits % bound; + } while (bits - value + (bound - 1) < 0L); + return value; + } + + private static float[] collectScores(IndexSearcher searcher, Query query, int maxDoc) + throws IOException { + int topN = Math.min(maxDoc, 10000); + TopDocs topDocs = searcher.search(query, topN); + float[] scores = new float[topDocs.scoreDocs.length]; + for (int i = 0; i < topDocs.scoreDocs.length; i++) { + scores[i] = topDocs.scoreDocs[i].score; + } + return scores; + } +} \ No newline at end of file diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/lucene/search/BayesianScoreQuery.java b/standalone-metastore/metastore-search/src/main/java/org/apache/lucene/search/BayesianScoreQuery.java new file mode 100644 index 000000000000..69036058ae1a --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/lucene/search/BayesianScoreQuery.java @@ -0,0 +1,295 @@ +/* + * 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.lucene.search; + +import java.io.IOException; +import java.util.Objects; + +import org.apache.lucene.index.LeafReaderContext; + +/** + * A query wrapper that transforms the inner query's score into a calibrated probability via sigmoid + * calibration: {@code P = sigmoid(alpha * (score - beta))}. + * + *

This implements the query-level Bayesian transform from "Bayesian BM25": the inner query + * (typically a multi-term BooleanQuery with BM25Similarity) produces a raw score, and this wrapper + * maps it to a probability in (0, 1) suitable for combination with other probability signals via + * {@link LogOddsFusionQuery}. + * + *

The alpha parameter controls the sigmoid steepness (score sensitivity), and beta controls the + * midpoint (decision boundary). These can be set manually or estimated from the score distribution + * via {@link BayesianScoreEstimator}. + * + *

An optional base rate encodes the corpus-level prior probability that a random document is + * relevant to a random query. When set, the posterior is computed in log-odds space: {@code + * sigmoid(alpha * (score - beta) + logit(baseRate))}. This shifts scores down for rare-relevance + * corpora, improving calibration. + * + * @lucene.experimental + */ +public final class BayesianScoreQuery extends Query { + + private final Query query; + private final float alpha; + private final float beta; + private final float baseRate; + private final float logitBaseRate; + + /** + * Creates a BayesianScoreQuery with base rate. + * + * @param query the inner query whose scores will be transformed + * @param alpha sigmoid steepness (must be positive and finite) + * @param beta sigmoid midpoint (must be finite) + * @param baseRate corpus-level relevance prior in (0, 1), or 0 to disable. When positive, adds + * logit(baseRate) to the log-odds before sigmoid, shifting scores to account for the rarity + * of relevant documents. + */ + public BayesianScoreQuery(Query query, float alpha, float beta, float baseRate) { + this.query = Objects.requireNonNull(query); + if (Float.isFinite(alpha) == false || alpha <= 0) { + throw new IllegalArgumentException("alpha must be a positive finite value, got " + alpha); + } + if (Float.isFinite(beta) == false) { + throw new IllegalArgumentException("beta must be a finite value, got " + beta); + } + if (baseRate < 0 || baseRate >= 1) { + throw new IllegalArgumentException("baseRate must be in [0, 1), got " + baseRate); + } + this.alpha = alpha; + this.beta = beta; + this.baseRate = baseRate; + if (baseRate > 0) { + this.logitBaseRate = (float) Math.log(baseRate / (1.0 - baseRate)); + } else { + this.logitBaseRate = 0f; + } + } + + /** + * Creates a BayesianScoreQuery without base rate. + * + * @param query the inner query whose scores will be transformed + * @param alpha sigmoid steepness (must be positive and finite) + * @param beta sigmoid midpoint (must be finite) + */ + public BayesianScoreQuery(Query query, float alpha, float beta) { + this(query, alpha, beta, 0f); + } + + /** Returns the wrapped query. */ + public Query getQuery() { + return query; + } + + /** Returns the sigmoid steepness parameter. */ + public float getAlpha() { + return alpha; + } + + /** Returns the sigmoid midpoint parameter. */ + public float getBeta() { + return beta; + } + + /** Returns the base rate, or 0 if not set. */ + public float getBaseRate() { + return baseRate; + } + + static float sigmoid(float x) { + if (x >= 0) { + return (float) (1.0 / (1.0 + Math.exp(-x))); + } else { + double expX = Math.exp(x); + return (float) (expX / (1.0 + expX)); + } + } + + @Override + public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) + throws IOException { + Weight innerWeight = query.createWeight(searcher, scoreMode, boost); + if (scoreMode.needsScores() == false) { + return innerWeight; + } + return new BayesianScoreWeight(this, innerWeight); + } + + @Override + public Query rewrite(IndexSearcher indexSearcher) throws IOException { + Query rewritten = query.rewrite(indexSearcher); + if (rewritten instanceof MatchNoDocsQuery) { + return rewritten; + } + if (rewritten != query) { + return new BayesianScoreQuery(rewritten, alpha, beta, baseRate); + } + return super.rewrite(indexSearcher); + } + + @Override + public void visit(QueryVisitor visitor) { + query.visit(visitor.getSubVisitor(BooleanClause.Occur.MUST, this)); + } + + @Override + public String toString(String field) { + String base = "BayesianScore(" + query.toString(field) + ", alpha=" + alpha + ", beta=" + beta; + if (baseRate > 0) { + base += ", baseRate=" + baseRate; + } + return base + ")"; + } + + @Override + public boolean equals(Object other) { + return sameClassAs(other) && equalsTo(getClass().cast(other)); + } + + private boolean equalsTo(BayesianScoreQuery other) { + return query.equals(other.query) + && Float.floatToIntBits(alpha) == Float.floatToIntBits(other.alpha) + && Float.floatToIntBits(beta) == Float.floatToIntBits(other.beta) + && Float.floatToIntBits(baseRate) == Float.floatToIntBits(other.baseRate); + } + + @Override + public int hashCode() { + int h = classHash(); + h = 31 * h + query.hashCode(); + h = 31 * h + Float.floatToIntBits(alpha); + h = 31 * h + Float.floatToIntBits(beta); + h = 31 * h + Float.floatToIntBits(baseRate); + return h; + } + + private class BayesianScoreWeight extends Weight { + private final Weight innerWeight; + + BayesianScoreWeight(Query query, Weight innerWeight) { + super(query); + this.innerWeight = innerWeight; + } + + @Override + public Matches matches(LeafReaderContext context, int doc) throws IOException { + return innerWeight.matches(context, doc); + } + + @Override + public Explanation explain(LeafReaderContext context, int doc) throws IOException { + Explanation innerExpl = innerWeight.explain(context, doc); + if (innerExpl.isMatch() == false) { + return innerExpl; + } + float innerScore = innerExpl.getValue().floatValue(); + float logOdds = alpha * (innerScore - beta) + logitBaseRate; + float transformed = sigmoid(logOdds); + if (baseRate > 0) { + return Explanation.match( + transformed, + "sigmoid calibration with base rate, computed as" + + " sigmoid(alpha * (score - beta) + logit(baseRate)) from:", + innerExpl, + Explanation.match(alpha, "alpha, sigmoid steepness"), + Explanation.match(beta, "beta, sigmoid midpoint"), + Explanation.match(baseRate, "baseRate, corpus-level relevance prior")); + } + return Explanation.match( + transformed, + "sigmoid calibration, computed as sigmoid(alpha * (score - beta)) from:", + innerExpl, + Explanation.match(alpha, "alpha, sigmoid steepness"), + Explanation.match(beta, "beta, sigmoid midpoint")); + } + + @Override + public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOException { + ScorerSupplier innerSupplier = innerWeight.scorerSupplier(context); + if (innerSupplier == null) { + return null; + } + return new ScorerSupplier() { + @Override + public Scorer get(long leadCost) throws IOException { + Scorer innerScorer = innerSupplier.get(leadCost); + return new BayesianScoreScorer(innerScorer); + } + + @Override + public long cost() { + return innerSupplier.cost(); + } + + @Override + public void setTopLevelScoringClause() throws IOException { + innerSupplier.setTopLevelScoringClause(); + } + }; + } + + @Override + public int count(LeafReaderContext context) throws IOException { + return innerWeight.count(context); + } + + @Override + public boolean isCacheable(LeafReaderContext ctx) { + return innerWeight.isCacheable(ctx); + } + } + + private class BayesianScoreScorer extends FilterScorer { + + BayesianScoreScorer(Scorer in) { + super(in); + } + + @Override + public float score() throws IOException { + float innerScore = in.score(); + return sigmoid(alpha * (innerScore - beta) + logitBaseRate); + } + + @Override + public int advanceShallow(int target) throws IOException { + return in.advanceShallow(target); + } + + @Override + public float getMaxScore(int upTo) throws IOException { + float innerMax = in.getMaxScore(upTo); + // sigmoid is monotone, so max(sigmoid(f(x))) = sigmoid(max(f(x))) + return sigmoid(alpha * (innerMax - beta) + logitBaseRate); + } + + @Override + public void setMinCompetitiveScore(float minScore) throws IOException { + // Invert the sigmoid to get the minimum inner score needed: + // minScore = sigmoid(alpha * (innerScore - beta) + logitBaseRate) + // => alpha * (innerScore - beta) + logitBaseRate = logit(minScore) + // => innerScore = (logit(minScore) - logitBaseRate) / alpha + beta + if (minScore > 0f && minScore < 1f) { + float clamped = Math.max(1e-7f, Math.min(1f - 1e-7f, minScore)); + float logitMin = (float) Math.log(clamped / (1f - clamped)); + float innerMin = (logitMin - logitBaseRate) / alpha + beta; + in.setMinCompetitiveScore(Math.max(0f, innerMin)); + } + } + } +} diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/lucene/search/LogOddsFusionQuery.java b/standalone-metastore/metastore-search/src/main/java/org/apache/lucene/search/LogOddsFusionQuery.java new file mode 100644 index 000000000000..0b01afc15ec0 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/lucene/search/LogOddsFusionQuery.java @@ -0,0 +1,479 @@ +/* + * 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.lucene.search; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; +import java.util.Objects; +import java.util.stream.Collectors; + +import org.apache.lucene.index.LeafReaderContext; + +/** + * A query that combines sub-query probability scores via log-odds fusion. Sub-queries are expected + * to produce scores in (0, 1) representing probabilities (e.g., from {@link BayesianScoreQuery} + * wrapping a BM25 query, or KNN cosine similarity). + * + *

The combination formula resolves the shrinkage problem of naive probabilistic AND by: + * + *

    + *
  1. Converting each sub-score to log-odds: logit(p) = log(p / (1 - p)) + *
  2. Computing the mean log-odds across all clauses (non-matching contribute 0 = neutral) + *
  3. Applying multiplicative confidence scaling: meanLogit * n^alpha + *
  4. Converting back to probability via sigmoid + *
+ * + *

The alpha parameter controls the confidence scaling exponent. The default alpha=0.5 implements + * the sqrt(n) scaling law from "From Bayesian Inference to Neural Computation". + * + *

Optional per-signal weights enable weighted Log-OP (Logarithmic Opinion Pooling) where each + * signal's log-odds contribution is scaled by its reliability weight. Weights must be non-negative + * and sum to 1. When weights are provided, the scoring formula becomes: {@code sigmoid(n^alpha * + * sum(w_i * softplus(logit(p_i))))} instead of the uniform mean. + * + * @see LogOddsFusionScorer + * @lucene.experimental + */ +public final class LogOddsFusionQuery extends Query implements Iterable { + + private final Multiset clauses = new Multiset<>(); + private final List orderedClauses; + private final float alpha; + private final float[] signalWeights; + private final float[] logitMin; + private final float[] logitMax; + + /** + * Creates a new LogOddsFusionQuery with per-signal weights and optional logit normalization. + * + * @param clauses the sub-queries to combine + * @param alpha confidence scaling exponent (0.5 = sqrt(n) law) + * @param weights per-signal weights (must be non-negative, finite, and sum to 1.0), or null for + * uniform weighting + * @param logitMin per-signal logit lower bounds for normalization, or null to use softplus gating + * @param logitMax per-signal logit upper bounds for normalization, or null to use softplus gating + * @throws IllegalArgumentException if alpha is not in [0, 1], or weights/bounds are invalid + */ + public LogOddsFusionQuery( + Collection clauses, + float alpha, + float[] weights, + float[] logitMin, + float[] logitMax) { + Objects.requireNonNull(clauses, "Collection of Queries must not be null"); + if (Float.isNaN(alpha) || alpha < 0 || alpha > 1) { + throw new IllegalArgumentException("alpha must be in [0, 1], got " + alpha); + } + if (weights != null) { + if (weights.length != clauses.size()) { + throw new IllegalArgumentException( + "weights length " + weights.length + " must equal clauses size " + clauses.size()); + } + float sum = 0; + for (float w : weights) { + if (Float.isFinite(w) == false || w < 0) { + throw new IllegalArgumentException("weights must be non-negative and finite, got " + w); + } + sum += w; + } + if (Math.abs(sum - 1.0f) > 1e-3f) { + throw new IllegalArgumentException("weights must sum to 1.0, got " + sum); + } + this.signalWeights = weights.clone(); + } else { + this.signalWeights = null; + } + if (logitMin != null && logitMax != null) { + if (logitMin.length != clauses.size()) { + throw new IllegalArgumentException( + "logitMin length " + logitMin.length + " must equal clauses size " + clauses.size()); + } + if (logitMax.length != clauses.size()) { + throw new IllegalArgumentException( + "logitMax length " + logitMax.length + " must equal clauses size " + clauses.size()); + } + this.logitMin = logitMin.clone(); + this.logitMax = logitMax.clone(); + } else { + this.logitMin = null; + this.logitMax = null; + } + this.alpha = alpha; + this.clauses.addAll(clauses); + this.orderedClauses = new ArrayList<>(clauses); + } + + /** + * Creates a new LogOddsFusionQuery with per-signal weights (softplus gating, no normalization). + * + * @param clauses the sub-queries to combine + * @param alpha confidence scaling exponent (0.5 = sqrt(n) law) + * @param weights per-signal weights, or null for uniform weighting + * @throws IllegalArgumentException if alpha is not in [0, 1], or weights are invalid + */ + public LogOddsFusionQuery(Collection clauses, float alpha, float[] weights) { + this(clauses, alpha, weights, null, null); + } + + /** + * Creates a new LogOddsFusionQuery with uniform weighting and softplus gating. + * + * @param clauses the sub-queries to combine + * @param alpha confidence scaling exponent (0.5 = sqrt(n) law) + * @throws IllegalArgumentException if alpha is not in [0, 1] + */ + public LogOddsFusionQuery(Collection clauses, float alpha) { + this(clauses, alpha, null, null, null); + } + + /** + * Creates a new LogOddsFusionQuery with default alpha=0.5, uniform weighting, and softplus + * gating. + * + * @param clauses the sub-queries to combine + */ + public LogOddsFusionQuery(Collection clauses) { + this(clauses, 0.5f, null, null, null); + } + + @Override + public Iterator iterator() { + return getClauses().iterator(); + } + + /** Returns the clauses. */ + public Collection getClauses() { + return Collections.unmodifiableCollection(orderedClauses); + } + + /** Returns the alpha (confidence scaling exponent). */ + public float getAlpha() { + return alpha; + } + + /** + * Returns a copy of the per-signal weights, or null if uniform weighting is used. + * + *

When non-null, the i-th element is the weight for the i-th clause in the order returned by + * {@link #getClauses()}. + */ + public float[] getWeights() { + return signalWeights != null ? signalWeights.clone() : null; + } + + /** Weight for LogOddsFusionQuery. */ + protected class LogOddsFusionWeight extends Weight { + + protected final ArrayList weights = new ArrayList<>(); + private final ScoreMode scoreMode; + + public LogOddsFusionWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) + throws IOException { + super(LogOddsFusionQuery.this); + for (Query clauseQuery : orderedClauses) { + weights.add(searcher.createWeight(clauseQuery, scoreMode, boost)); + } + this.scoreMode = scoreMode; + } + + @Override + public Matches matches(LeafReaderContext context, int doc) throws IOException { + List mis = new ArrayList<>(); + for (Weight weight : weights) { + Matches mi = weight.matches(context, doc); + if (mi != null) { + mis.add(mi); + } + } + return MatchesUtils.fromSubMatches(mis); + } + + @Override + public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOException { + List scorerSuppliers = new ArrayList<>(); + List activeWeightsList = signalWeights != null ? new ArrayList<>() : null; + List activeLogitMinList = logitMin != null ? new ArrayList<>() : null; + List activeLogitMaxList = logitMax != null ? new ArrayList<>() : null; + + for (int i = 0; i < weights.size(); i++) { + ScorerSupplier ss = weights.get(i).scorerSupplier(context); + if (ss != null) { + scorerSuppliers.add(ss); + if (activeWeightsList != null) { + activeWeightsList.add(signalWeights[i]); + } + if (activeLogitMinList != null) { + activeLogitMinList.add(logitMin[i]); + activeLogitMaxList.add(logitMax[i]); + } + } + } + + if (scorerSuppliers.isEmpty()) { + return null; + } else if (scorerSuppliers.size() == 1) { + return scorerSuppliers.get(0); + } else { + final int totalClauses = clauses.size(); + final float[] activeWeights = toFloatArray(activeWeightsList); + final float[] activeMin = toFloatArray(activeLogitMinList); + final float[] activeMax = toFloatArray(activeLogitMaxList); + + return new ScorerSupplier() { + + private long cost = -1; + + @Override + public Scorer get(long leadCost) throws IOException { + List scorers = new ArrayList<>(); + for (ScorerSupplier ss : scorerSuppliers) { + scorers.add(ss.get(leadCost)); + } + return new LogOddsFusionScorer( + scorers, + totalClauses, + alpha, + activeWeights, + activeMin, + activeMax, + scoreMode, + leadCost); + } + + @Override + public long cost() { + if (cost == -1) { + long cost = 0; + for (ScorerSupplier ss : scorerSuppliers) { + cost += ss.cost(); + } + this.cost = cost; + } + return cost; + } + + @Override + public void setTopLevelScoringClause() throws IOException { + for (ScorerSupplier ss : scorerSuppliers) { + ss.setTopLevelScoringClause(); + } + } + }; + } + } + + @Override + public boolean isCacheable(LeafReaderContext ctx) { + if (weights.size() + > AbstractMultiTermQueryConstantScoreWrapper.BOOLEAN_REWRITE_TERM_COUNT_THRESHOLD) { + return false; + } + for (Weight w : weights) { + if (w.isCacheable(ctx) == false) return false; + } + return true; + } + + @Override + public Explanation explain(LeafReaderContext context, int doc) throws IOException { + boolean match = false; + List subsOnMatch = new ArrayList<>(); + List subsOnNoMatch = new ArrayList<>(); + double logitSum = 0; + int totalClauses = weights.size(); + + for (int i = 0; i < weights.size(); i++) { + Explanation e = weights.get(i).explain(context, doc); + if (e.isMatch()) { + match = true; + subsOnMatch.add(e); + float subScore = e.getValue().floatValue(); + float rawLogit = LogOddsFusionScorer.logit(subScore); + float gated; + if (logitMin != null) { + float range = logitMax[i] - logitMin[i]; + gated = range > 0 ? Math.clamp((rawLogit - logitMin[i]) / range, 0f, 1f) : 0.5f; + } else { + gated = LogOddsFusionScorer.softplus(rawLogit); + } + if (signalWeights != null) { + logitSum += signalWeights[i] * gated; + } else { + logitSum += gated; + } + } else if (match == false) { + subsOnNoMatch.add(e); + } + } + + if (match) { + float scalingFactor = (float) Math.pow(totalClauses, alpha); + float scaledLogit = 0f; + String description; + if (signalWeights != null) { + scaledLogit = (float) logitSum * scalingFactor; + description = + "weighted log-odds fusion, computed as sigmoid(weightedLogit * n^alpha) from:"; + } else { + scaledLogit = (float) (logitSum / totalClauses) * scalingFactor; + description = "log-odds fusion, computed as sigmoid(meanLogit * n^alpha) from:"; + } + float score = LogOddsFusionScorer.sigmoid(scaledLogit); + return Explanation.match(score, description, subsOnMatch); + } else { + return Explanation.noMatch("No matching clause", subsOnNoMatch); + } + } + } + + @Override + public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) + throws IOException { + return new LogOddsFusionWeight(searcher, scoreMode, boost); + } + + @Override + public Query rewrite(IndexSearcher indexSearcher) throws IOException { + if (clauses.isEmpty()) { + return new MatchNoDocsQuery("empty LogOddsFusionQuery"); + } + + if (clauses.size() == 1) { + return orderedClauses.get(0); + } + + boolean actuallyRewritten = false; + List rewrittenClauses = new ArrayList<>(); + List newWeights = signalWeights != null ? new ArrayList<>() : null; + List newLogitMin = logitMin != null ? new ArrayList<>() : null; + List newLogitMax = logitMax != null ? new ArrayList<>() : null; + + for (int i = 0; i < orderedClauses.size(); i++) { + Query sub = orderedClauses.get(i); + Query rewrittenSub = sub.rewrite(indexSearcher); + if (rewrittenSub != sub || sub.getClass() == MatchNoDocsQuery.class) { + actuallyRewritten = true; + } + if (rewrittenSub.getClass() != MatchNoDocsQuery.class) { + rewrittenClauses.add(rewrittenSub); + if (newWeights != null) { + newWeights.add(signalWeights[i]); + } + if (newLogitMin != null) { + newLogitMin.add(logitMin[i]); + newLogitMax.add(logitMax[i]); + } + } + } + + if (actuallyRewritten == false) { + return super.rewrite(indexSearcher); + } + if (rewrittenClauses.isEmpty()) { + return new MatchNoDocsQuery("empty LogOddsFusionQuery"); + } + if (rewrittenClauses.size() == 1) { + return rewrittenClauses.get(0); + } + + float[] filteredWeights = toFloatArray(newWeights); + if (filteredWeights != null) { + float sum = 0; + for (float w : filteredWeights) { + sum += w; + } + if (sum > 0) { + for (int i = 0; i < filteredWeights.length; i++) { + filteredWeights[i] /= sum; + } + } + } + + return new LogOddsFusionQuery( + rewrittenClauses, + alpha, + filteredWeights, + toFloatArray(newLogitMin), + toFloatArray(newLogitMax)); + } + + @Override + public void visit(QueryVisitor visitor) { + QueryVisitor v = visitor.getSubVisitor(BooleanClause.Occur.SHOULD, this); + for (Query q : clauses) { + q.visit(v); + } + } + + @Override + public String toString(String field) { + String base = + this.orderedClauses.stream() + .map( + subquery -> { + if (subquery instanceof BooleanQuery) { + return "(" + subquery.toString(field) + ")"; + } + return subquery.toString(field); + }) + .collect(Collectors.joining(" & ", "LogOdds(", ")^" + alpha)); + if (signalWeights != null) { + return base + " w=" + Arrays.toString(signalWeights); + } + return base; + } + + @Override + public boolean equals(Object other) { + return sameClassAs(other) && equalsTo(getClass().cast(other)); + } + + private boolean equalsTo(LogOddsFusionQuery other) { + return alpha == other.alpha + && Objects.equals(clauses, other.clauses) + && Arrays.equals(signalWeights, other.signalWeights) + && Arrays.equals(logitMin, other.logitMin) + && Arrays.equals(logitMax, other.logitMax); + } + + @Override + public int hashCode() { + int h = classHash(); + h = 31 * h + Float.floatToIntBits(alpha); + h = 31 * h + Objects.hashCode(clauses); + h = 31 * h + Arrays.hashCode(signalWeights); + h = 31 * h + Arrays.hashCode(logitMin); + h = 31 * h + Arrays.hashCode(logitMax); + return h; + } + + private static float[] toFloatArray(List list) { + if (list == null) { + return null; + } + float[] result = new float[list.size()]; + for (int i = 0; i < list.size(); i++) { + result[i] = list.get(i); + } + return result; + } +} \ No newline at end of file diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/lucene/search/LogOddsFusionScorer.java b/standalone-metastore/metastore-search/src/main/java/org/apache/lucene/search/LogOddsFusionScorer.java new file mode 100644 index 000000000000..5efa9d6deaf8 --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/lucene/search/LogOddsFusionScorer.java @@ -0,0 +1,219 @@ +/* + * 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.lucene.search; + +import java.io.IOException; +import java.util.IdentityHashMap; +import java.util.List; + +/** + * Scorer for {@link LogOddsFusionQuery}. Combines sub-scorer outputs (assumed to be probabilities + * in (0, 1)) via log-odds fusion with multiplicative confidence scaling. + * + *

The scoring formula is: + * + * 

+ *   gatedLogit  = softplus(logit(clamp(subScore)))   for each matching sub-scorer
+ *   logitSum    = sum of gatedLogit values
+ *   meanLogit   = logitSum / n   (n = total clause count, not just matching)
+ *   scaledLogit = meanLogit * pow(n, alpha)
+ *   score       = sigmoid(scaledLogit)
+ *
+ * + *

Softplus gating ({@code log(1 + exp(x))}) is applied to logit values before aggregation. This + * distinguishes "absence of evidence" (non-matching sub-scorer, contributes 0) from "evidence of + * absence" (matching sub-scorer with weak probability, contributes a small positive value). A + * matching sub-scorer always contributes more than a non-matching one, preserving the ordering + * among weak matches while ensuring that no match is ever penalized. + * + *

Non-matching sub-scorers contribute logit(0.5) = 0 (neutral evidence). + * + * @lucene.experimental + */ +final class LogOddsFusionScorer extends DisjunctionScorer { + private static final float CLAMP_MIN = 1e-7f; + private static final float CLAMP_MAX = 1f - 1e-7f; + + private final List subScorers; + private final int totalClauses; + private final float scalingFactor; + private final float[] signalWeights; + private final float[] logitMin; + private final float[] logitMax; + private final IdentityHashMap scorerIndexMap; + + private final DisjunctionScoreBlockBoundaryPropagator disjunctionBlockPropagator; + + /** + * Creates a new LogOddsFusionScorer. + * + * @param subScorers the sub scorers to combine + * @param totalClauses the total number of clauses (including non-matching) + * @param alpha confidence scaling exponent (0.5 = sqrt(n) law) + * @param signalWeights per-signal weights parallel to subScorers (null for uniform weighting). + * When provided, the scoring formula uses weighted sum instead of mean. Weights must be + * non-negative and should sum to 1. + * @param logitMin per-signal logit lower bounds for normalization (null to use softplus gating). + * When provided together with logitMax, logit values are normalized to [0, 1] instead of + * applying softplus. This ensures non-negative contributions while preserving learned signal + * scale calibration. + * @param logitMax per-signal logit upper bounds for normalization (null to use softplus gating) + * @param scoreMode the score mode + * @param leadCost the lead cost for iteration + */ + LogOddsFusionScorer( + List subScorers, + int totalClauses, + float alpha, + float[] signalWeights, + float[] logitMin, + float[] logitMax, + ScoreMode scoreMode, + long leadCost) + throws IOException { + super(subScorers, scoreMode, leadCost); + this.subScorers = subScorers; + this.totalClauses = totalClauses; + this.scalingFactor = (float) Math.pow(totalClauses, alpha); + this.signalWeights = signalWeights; + this.logitMin = logitMin; + this.logitMax = logitMax; + if (signalWeights != null) { + this.scorerIndexMap = new IdentityHashMap<>(subScorers.size()); + for (int i = 0; i < subScorers.size(); i++) { + this.scorerIndexMap.put(subScorers.get(i), i); + } + } else { + this.scorerIndexMap = null; + } + if (scoreMode == ScoreMode.TOP_SCORES) { + this.disjunctionBlockPropagator = new DisjunctionScoreBlockBoundaryPropagator(subScorers); + } else { + this.disjunctionBlockPropagator = null; + } + } + + static float clampProbability(float p) { + return Math.clamp(p, CLAMP_MIN, CLAMP_MAX); + } + + static float logit(float p) { + float clamped = clampProbability(p); + return (float) Math.log(clamped / (1f - clamped)); + } + + static float sigmoid(float x) { + if (x >= 0) { + return (float) (1.0 / (1.0 + Math.exp(-x))); + } else { + double expX = Math.exp(x); + return (float) (expX / (1.0 + expX)); + } + } + + /** + * Softplus function: log(1 + exp(x)). Always positive, smooth approximation of ReLU. For large + * positive x, approaches x. For large negative x, approaches 0 from above. At x=0, returns log(2) + * ~ 0.693. + * + *

Uses a numerically stable formulation: for x > 20, softplus(x) ~ x. + */ + static float softplus(float x) { + if (x > 20f) { + return x; + } + return (float) Math.log1p(Math.exp(x)); + } + + /** Applies gating to a logit value: normalization if bounds are set, softplus otherwise. */ + private float gateLogit(float rawLogit, int signalIndex) { + if (logitMin != null) { + float range = logitMax[signalIndex] - logitMin[signalIndex]; + if (range > 0) { + return Math.clamp((rawLogit - logitMin[signalIndex]) / range, 0f, 1f); + } + return 0.5f; + } + return softplus(rawLogit); + } + + @Override + protected float score(DisiWrapper topList) throws IOException { + double logitSum = 0; + for (DisiWrapper w = topList; w != null; w = w.next) { + float subScore = w.scorable.score(); + int idx = scorerIndexMap != null ? scorerIndexMap.get(w.scorer) : -1; + float gated = gateLogit(logit(subScore), idx >= 0 ? idx : 0); + if (scorerIndexMap != null) { + logitSum += signalWeights[idx] * gated; + } else { + logitSum += gated; + } + } + // Non-matching sub-scorers contribute 0. + // With weights: sum(w_i * gated_i) already accounts for the 1/n factor. + // Without weights: divide by totalClauses to compute the mean. + float scaledLogit = 0f; + if (signalWeights != null) { + scaledLogit = (float) logitSum * scalingFactor; + } else { + scaledLogit = (float) (logitSum / totalClauses) * scalingFactor; + } + return sigmoid(scaledLogit); + } + + @Override + public int advanceShallow(int target) throws IOException { + if (disjunctionBlockPropagator != null) { + return disjunctionBlockPropagator.advanceShallow(target); + } + return super.advanceShallow(target); + } + + @Override + public float getMaxScore(int upTo) throws IOException { + // Safe upper bound: gateLogit is monotone in p (both softplus and normalize are monotone), + // weights are non-negative, sum of upper bounds >= sum of actuals, and sigmoid is monotone. + double maxLogitSum = 0; + for (int i = 0; i < subScorers.size(); i++) { + Scorer scorer = subScorers.get(i); + if (scorer.docID() <= upTo) { + float maxSubScore = scorer.getMaxScore(upTo); + float gated = gateLogit(logit(maxSubScore), i); + if (signalWeights != null) { + maxLogitSum += signalWeights[i] * gated; + } else { + maxLogitSum += gated; + } + } + } + float scaledLogit = 0f; + if (signalWeights != null) { + scaledLogit = (float) maxLogitSum * scalingFactor; + } else { + scaledLogit = (float) (maxLogitSum / totalClauses) * scalingFactor; + } + return sigmoid(scaledLogit); + } + + @Override + public void setMinCompetitiveScore(float minScore) throws IOException { + if (disjunctionBlockPropagator != null) { + disjunctionBlockPropagator.setMinCompetitiveScore(minScore); + } + } +} \ No newline at end of file diff --git a/standalone-metastore/metastore-search/src/main/java/org/apache/lucene/search/package-info.java b/standalone-metastore/metastore-search/src/main/java/org/apache/lucene/search/package-info.java new file mode 100644 index 000000000000..a5b988f405aa --- /dev/null +++ b/standalone-metastore/metastore-search/src/main/java/org/apache/lucene/search/package-info.java @@ -0,0 +1,21 @@ +/* + * 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. + */ + +/** + * Classes in this package are copied from Lucene codebase. + */ +package org.apache.lucene.search; \ No newline at end of file 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..4fccd37ae821 --- /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(IndexStateConfig.REMOTE_URI, "file:///tmp/backup"); + conf.setBoolean(IndexStateConfig.MEMORY, true); + IndexStateConfig config = new IndexStateConfig(conf, "hive_tables"); + assertTrue(config.hasRemote()); + assertTrue(config.isDistributed()); + assertTrue(config.useMemory()); + } + + @Test + public void validateRemoteUriRejectsInvalidValue() { + assertThrows(IndexIOException.class, () -> IndexStateConfig.validateRemoteUri("://bad")); + } + + @Test + public void localPathDefaultsToWorkdir() { + IndexStateConfig config = new IndexStateConfig(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/TestSearchConfig.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/config/TestSearchConfig.java new file mode 100644 index 000000000000..c4794c64a631 --- /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.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() { + 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_MATCH_WEIGHT_DEFAULT, config.getHybridMatchWeight(), 0.001f); + assertEquals(SearchConfig.HYBRID_SEMANTIC_WEIGHT_DEFAULT, config.getHybridSemanticWeight(), 0.001f); + } + + @Test + public void readsConfiguredValues() { + Configuration conf = new Configuration(false); + conf.setInt(SearchConfig.DEFAULT_LIMIT, 50); + conf.setFloat(SearchConfig.HYBRID_MATCH_WEIGHT, 0.8f); + 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..44c944b2f2e0 --- /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.IndexStateConfig; +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(IndexStateConfig.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..1d35be82409c --- /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.IndexStateConfig; +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(IndexStateConfig.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..e07205b80554 --- /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.IndexStateConfig; +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(IndexStateConfig.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..1078a830b25c --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/index/TestIndexerSearchIntegration.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; + +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.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.get("_id").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).get("_id").toString().contains("sales.revenue")); + assertTrue(hits.get(1).get("_id").toString().contains("sales.line_items")); + assertTrue(hits.get(2).get("_id").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).get("_id").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..d71d1e114243 --- /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.encode(EmbedModel.TaskType.QUERY, "sales"); + float[] second = model.encode(EmbedModel.TaskType.QUERY, "sales"); + assertArrayEquals(first, second, 0.0001f); + + float[] document = model.encode(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/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..ae3dfe1ece24 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/RealMetastoreSearchSession.java @@ -0,0 +1,252 @@ +/* + * 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.IndexStateConfig; +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.SearchInternal; +import org.apache.hive.search.search.SearchReqResp; +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(IndexStateConfig.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).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).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 List> searchMatch(String text, int limit) throws Exception { + refreshSearcher(); + try (SearchInternal searchIO = new SearchInternal( + searcherManager, indexManager, modelRegistry, searchConfig, bayesianParameters)) { + return searchIO.search(SearchReqResp.Request.validated( + Map.of("table_keyword", text), + List.of(MetastoreTableMapper.FIELD_TABLE, MetastoreTableMapper.FIELD_COMMENT), + limit, + null, + null)).hits(); + } + } + + public List> searchHybrid(String queryText, int limit) throws Exception { + refreshSearcher(); + Map hybridBody = Map.of( + "field", MetastoreTableMapper.FIELD_SEARCH_TEXT, + "query", queryText); + try (SearchInternal searchIO = new SearchInternal( + searcherManager, indexManager, modelRegistry, searchConfig, bayesianParameters)) { + return searchIO.search(SearchReqResp.Request.validated( + Map.of("hybrid", hybridBody), + List.of(MetastoreTableMapper.FIELD_TABLE, MetastoreTableMapper.FIELD_COMMENT), + limit, + null, + null)).hits(); + } + } + + 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..279216b1fc49 --- /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.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.mapping.TableDocument; +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.BOOTSTRAP_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(org.apache.hive.search.config.IndexStateConfig.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..9d50c6a5d352 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestMetastoreEventHandlerIntegration.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.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 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..4b7b5f268b8f --- /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.IndexStateConfig; +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(IndexStateConfig.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..b31c842d89a8 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestMetastoreTableMapper.java @@ -0,0 +1,310 @@ +/* + * 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.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); + assertTrue(searchText.contains("orders")); + assertTrue(searchText.contains("daily orders")); + assertTrue(searchText.contains("id order id")); + assertFalse(searchText.contains("amount")); + 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 searchTextIncludesOnlyCommentedColumns() 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); + String searchText = fieldValue(document, MetastoreTableMapper.FIELD_SEARCH_TEXT); + String storedColumns = fieldValue(document, MetastoreTableMapper.FIELD_COLUMNS); + + assertTrue(searchText.contains("id order id")); + assertTrue(searchText.contains("status fulfillment status")); + assertFalse(searchText.contains("amount")); + assertTrue(storedColumns.contains("amount double")); + } + + @Test + public void searchTextCapsWideTables() 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 maxCols = MetastoreTableMapper.MAX_SEARCH_COLUMNS; + for (int i = 0; i < maxCols + 5; 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 storedColumns = fieldValue(document, MetastoreTableMapper.FIELD_COLUMNS); + + assertTrue(searchText.contains("col0 comment 0")); + assertTrue(searchText.contains("col" + (maxCols - 1) + " comment " + (maxCols - 1))); + assertFalse(searchText.contains("; col" + maxCols + " comment " + maxCols + ";")); + assertFalse(searchText.endsWith("; col" + maxCols + " comment " + maxCols)); + assertTrue(searchText.contains("(+5 more)")); + assertTrue(storedColumns.contains("col" + (maxCols + 4) + " string comment " + (maxCols + 4))); + assertFalse(storedColumns.contains("(+5 more)")); + } + + @Test + public void embedDocumentsPreservesLexicalFields() throws Exception { + Configuration conf = new Configuration(false); + conf.setBoolean(org.apache.hive.search.config.IndexStateConfig.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")); + } + } + + 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..69a07a8a6c95 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestRealMetastoreSearchFailureRecovery.java @@ -0,0 +1,115 @@ +/* + * 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.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 java.util.List; +import java.util.Map; + +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(); + List> hits = restored.searchMatch("sales", 5); + assertFalse(hits.isEmpty()); + assertTrue(hits.stream().anyMatch(hit -> hit.get("_id").toString().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).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).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..36c13ab05229 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/metastore/TestRealMetastoreSearchHybrid.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.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 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); + List> hits = session.searchMatch("sales", 5); + assertFalse(hits.isEmpty()); + assertTrue(hits.stream().anyMatch(hit -> hit.get("_id").toString().contains("orders"))); + assertTrue(session.searchMatch("parts", 5).stream() + .anyMatch(hit -> hit.get("_id").toString().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); + List> hits = session.searchHybrid("revenue analytics", 5); + assertFalse(hits.isEmpty()); + assertTrue(hits.stream().anyMatch(hit -> hit.get("_id").toString().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); + List> keywordHits = session.searchMatch("sales", 5); + assertFalse(keywordHits.isEmpty()); + assertTrue(keywordHits.get(0).get("_id").toString().contains("orders")); + + List> hybridHits = session.searchHybrid("revenue analytics", 5); + assertFalse(hybridHits.isEmpty()); + assertTrue(hybridHits.stream().anyMatch(hit -> hit.get("_id").toString().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..85ead8c5347d --- /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).isEmpty()); + assertFalse(session.searchMatch("customers", 10).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).isEmpty()); + + metastore.dropTable(db, "orders"); + session.waitUntilNotSearchable("sales", 5); + assertTrue(session.searchMatch("sales", 5).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).isEmpty()); + + metastore.dropTable(db, "orders"); + metastore.createTable(db, "orders", "revenue analytics dashboard"); + session.waitUntilSearchable("revenue", 5); + assertFalse(session.searchMatch("revenue", 5).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..a454493dd276 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/search/InMemorySearchFixture.java @@ -0,0 +1,173 @@ +/* + * 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.IndexStateConfig; +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(IndexStateConfig.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(Map.of("table_keyword", text), limit); + } + + public List> search(Map query, 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)) { + SearchReqResp.Response response = searchIO.search(SearchReqResp.Request.validated( + query, + List.of(MetastoreTableMapper.FIELD_TABLE, MetastoreTableMapper.FIELD_COMMENT), + limit, + null, + null)); + return response.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..59ad03002dba --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/search/TestHybridSearch.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.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.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 TestHybridSearch { + + private static IndexMapping hybridMapping() { + Configuration conf = new Configuration(false); + return MetastoreSchemas.defaultHiveTablesMapping("hive_tables", "bge-small", conf); + } + + @Test + public void parseStringUsesSoleHybridField() throws Exception { + HybridSearch.ParsedHybridQuery parsed = + HybridSearch.parse("sales revenue", hybridMapping()); + assertEquals(MetastoreTableMapper.FIELD_SEARCH_TEXT, parsed.field()); + assertEquals("sales revenue", parsed.queryText()); + } + + @Test + public void parseFieldAndQueryObject() throws Exception { + Map body = Map.of( + "field", MetastoreTableMapper.FIELD_SEARCH_TEXT, + "query", "orders"); + HybridSearch.ParsedHybridQuery parsed = HybridSearch.parse(body, hybridMapping()); + assertEquals(MetastoreTableMapper.FIELD_SEARCH_TEXT, parsed.field()); + assertEquals("orders", parsed.queryText()); + } + + @Test + public void parseSingleFieldMap() throws Exception { + Map body = Map.of(MetastoreTableMapper.FIELD_SEARCH_TEXT, "inventory"); + HybridSearch.ParsedHybridQuery parsed = HybridSearch.parse(body, hybridMapping()); + assertEquals("inventory", parsed.queryText()); + } + + @Test + public void parseCustomWeights() throws Exception { + Map body = Map.of( + "field", MetastoreTableMapper.FIELD_SEARCH_TEXT, + "query", "metrics", + "match_weight", 0.7, + "semantic_weight", 0.3); + HybridSearch.ParsedHybridQuery parsed = HybridSearch.parse(body, hybridMapping()); + assertEquals(0.7f, parsed.matchWeight(), 0.001f); + assertEquals(0.3f, parsed.semanticWeight(), 0.001f); + } + + @Test + public void parseRejectsNonHybridField() { + Map body = Map.of( + "field", MetastoreTableMapper.FIELD_TABLE, + "query", "t1"); + assertThrows(SearchException.class, () -> HybridSearch.parse(body, hybridMapping())); + } + + @Test + public void parseRequiresFieldWhenMultipleHybridFieldsExist() { + 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, () -> HybridSearch.parse("query", mapping)); + } + + @Test + public void toFusionRequestUsesDefaultsFromSearchConfig() { + HybridSearch.ParsedHybridQuery parsed = + new HybridSearch.ParsedHybridQuery(MetastoreTableMapper.FIELD_SEARCH_TEXT, "sales", null, null); + SearchConfig searchConfig = new SearchConfig(new Configuration(false)); + SearchInternal.FusionRequest request = + HybridSearch.toFusionRequest(parsed, 20, searchConfig); + + assertEquals(20, request.size()); + assertEquals(2, request.retrievers().size()); + assertEquals("match", request.retrievers().get(0).name()); + assertEquals( + "sales", + request.retrievers().get(0).query().get("table_keyword")); + 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/TestSearchQuery.java b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/search/TestSearchQuery.java new file mode 100644 index 000000000000..1d0b78932fc7 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/java/org/apache/hive/search/search/TestSearchQuery.java @@ -0,0 +1,67 @@ +/* + * 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.common.TableName; +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; + +@Category(MetastoreUnitTest.class) +public class TestSearchQuery { + + @Test + public void ofTextDefaultsToHybridMode() throws Exception { + SearchQuery query = SearchQuery.of("sales"); + assertEquals("sales", query.queryText()); + assertEquals(SearchReqResp.Mode.HYBRID, query.mode()); + assertEquals(0, query.limit()); + } + + @Test + public void ofTextWithModeAndLimit() throws Exception { + SearchQuery query = SearchQuery.of("sales", SearchReqResp.Mode.SEMANTIC, 25); + assertEquals(SearchReqResp.Mode.SEMANTIC, query.mode()); + assertEquals(25, query.limit()); + } + + @Test + public void ofTableNameUsesKeywordMode() throws Exception { + TableName tableName = new TableName("hive", "default", "orders"); + SearchQuery query = SearchQuery.of(tableName); + assertEquals("orders", query.queryText()); + assertEquals(SearchReqResp.Mode.KEYWORD, 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", SearchReqResp.Mode.HYBRID, -1)); + } +} 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..1e49dffa9267 --- /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[] encode(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[][] encodeBatch(TaskType task, String[] texts) throws IndexException { + encodeBatchCalls.incrementAndGet(); + return Arrays.stream(texts).map(text -> encode(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/metastore-search/src/test/resources/log4j2.properties b/standalone-metastore/metastore-search/src/test/resources/log4j2.properties new file mode 100644 index 000000000000..2e185e152d23 --- /dev/null +++ b/standalone-metastore/metastore-search/src/test/resources/log4j2.properties @@ -0,0 +1,27 @@ +# 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. + +name = PropertiesConfig + +appender.console.type = Console +appender.console.name = STDOUT +appender.console.layout.type = PatternLayout +appender.console.layout.pattern = [%-5level] %d{yyyy-MM-dd HH:mm:ss.SSS} [%t] %c{1} - %msg%n + +rootLogger.level = info +rootLogger.appenderRefs = stdout +rootLogger.appenderRef.stdout.ref = STDOUT 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