SynonymV2GraphFilterFactory and NrtsearchSynonymParser

src/main/java/com/yelp/nrtsearch/server/luceneserver/analysis/NrtsearchSynonymParser.java

@@ -0,0 +1,137 @@
+/*
+ * Copyright 2024 Yelp Inc.
+ *
+ * Licensed 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 com.yelp.nrtsearch.server.luceneserver.analysis;
+
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.io.Reader;
+import java.text.ParseException;
+import java.util.ArrayList;
+import org.apache.lucene.analysis.Analyzer;
+import org.apache.lucene.analysis.synonym.SynonymMap;
+import org.apache.lucene.util.CharsRef;
+import org.apache.lucene.util.CharsRefBuilder;
+
+class NrtsearchSynonymParser extends SynonymMap.Parser {
+  private final boolean expand;
+  private static final String SYNONYMS_SEPARATOR = "\\|";
+  private static final String SYNONYM_MAPPING_SEPARATOR = ",";
+
+  public NrtsearchSynonymParser(boolean dedup, boolean expand, Analyzer analyzer) {
+    super(dedup, analyzer);
+    this.expand = expand;
+  }
+
+  @Override
+  public void parse(Reader mappings) throws IOException, ParseException {
+    BufferedReader bufferedReader = new BufferedReader(mappings);
+    String line;
+    while ((line = bufferedReader.readLine()) != null) {
+      // Splitting the mappings based on "|"
+      String[] synonyms = line.split(SYNONYMS_SEPARATOR);
+      this.addInternal(synonyms);
+    }
+  }
+
+  public void addInternal(String[] synonyms) throws IOException {
+    String[] inputStrings;
+    CharsRef[] inputs;
+    int i;
+
+    for (String synonym : synonyms) {
+      inputStrings = split(synonym, SYNONYM_MAPPING_SEPARATOR);
+
+      if (inputStrings.length != 2) {
+        throw new IllegalArgumentException("synonym mapping is invalid for " + synonym);
+      }
+      inputs = new CharsRef[inputStrings.length];
+
+      for (i = 0; i < inputs.length; ++i) {
+        inputs[i] = this.analyze(this.unescape(inputStrings[i]).trim(), new CharsRefBuilder());
+      }
+
+      if (!this.expand) {
+        for (i = 0; i < inputs.length; ++i) {
+          this.add(inputs[i], inputs[0], false);
+        }
+      } else {
+        for (i = 0; i < inputs.length; ++i) {
+          for (int j = 0; j < inputs.length; ++j) {
+            if (i != j) {
+              this.add(inputs[i], inputs[j], true);
+            }
+          }
+        }
+      }
+    }
+  }
+
+  private static String[] split(String s, String separator) {
+    ArrayList<String> list = new ArrayList(2);
+    StringBuilder sb = new StringBuilder();
+    int pos = 0;
+    int end = s.length();
+
+    while (pos < end) {
+      if (s.startsWith(separator, pos)) {
+        if (sb.length() > 0) {
+          list.add(sb.toString());
+          sb = new StringBuilder();
+        }
+
+        pos += separator.length();
+      } else {
+        char ch = s.charAt(pos++);
+        if (ch == '\\') {
+          sb.append(ch);
+          if (pos >= end) {
+            break;
+          }
+
+          ch = s.charAt(pos++);
+        }
+
+        sb.append(ch);
+      }
+    }
+
+    if (sb.length() > 0) {
+      list.add(sb.toString());
+    }
+
+    return list.toArray(new String[list.size()]);
+  }
+
+  private String unescape(String s) {
+    if (s.indexOf("\\") < 0) {
+      return s;
+    } else {
+      StringBuilder sb = new StringBuilder();
+
+      for (int i = 0; i < s.length(); ++i) {
+        char ch = s.charAt(i);
+        if (ch == '\\' && i < s.length() - 1) {
+          ++i;
+          sb.append(s.charAt(i));
+        } else {
+          sb.append(ch);
+        }
+      }
+
+      return sb.toString();
+    }
+  }
+}

src/main/java/com/yelp/nrtsearch/server/luceneserver/analysis/SynonymV2GraphFilterFactory.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright 2024 Yelp Inc.
+ *
+ * Licensed 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 com.yelp.nrtsearch.server.luceneserver.analysis;
+
+import java.io.IOException;
+import java.io.StringReader;
+import java.text.ParseException;
+import java.util.Map;
+import org.apache.lucene.analysis.Analyzer;
+import org.apache.lucene.analysis.TokenStream;
+import org.apache.lucene.analysis.synonym.SolrSynonymParser;
+import org.apache.lucene.analysis.synonym.SynonymGraphFilter;
+import org.apache.lucene.analysis.synonym.SynonymMap;
+import org.apache.lucene.analysis.synonym.WordnetSynonymParser;
+import org.apache.lucene.analysis.util.TokenFilterFactory;
+
+public class SynonymV2GraphFilterFactory extends TokenFilterFactory {
+  public static final String MAPPINGS = "mappings";
+  public final boolean ignoreCase;
+  protected SynonymMap synonymMap;
+
+  public SynonymV2GraphFilterFactory(Map<String, String> args, Analyzer analyzer)
+      throws IOException, ParseException {
+    super(args);
+    this.ignoreCase = getBoolean(args, "ignoreCase", false);
+    boolean expand = getBoolean(args, "expand", true);
+    String parserFormat = args.get("parserFormat");
+    String synonymMappings = args.get(MAPPINGS);
+    if (synonymMappings == null) {
+      throw new IllegalArgumentException("Synonym mappings must be specified");
+    }
+    if (parserFormat == null) {
+      throw new IllegalArgumentException("Parser format must be specified");
+    }
+    synonymMap = loadSynonymsFromString(synonymMappings, parserFormat, expand, analyzer);
+  }
+
+  @Override
+  public TokenStream create(TokenStream input) {
+    return new SynonymGraphFilter(input, synonymMap, ignoreCase);
+  }
+
+  public SynonymMap loadSynonymsFromString(
+      String synonymMappings, String parserFormat, boolean expand, Analyzer analyzer)
+      throws IOException, ParseException {
+    SynonymMap.Parser parser;
+
+    if (parserFormat.equals("solr")) {
+      parser = new SolrSynonymParser(true, expand, analyzer);
+    } else if (parserFormat.equals("wordnet")) {
+      parser = new WordnetSynonymParser(true, expand, analyzer);
+    } else if (parserFormat.equals("nrtsearch")) {
+      parser = new NrtsearchSynonymParser(true, expand, analyzer);
+    } else {
+      throw new IllegalArgumentException(
+          "The parser format: "
+              + parserFormat
+              + " is not valid. It should be solr, wordnet or nrtsearch");
+    }
+    parser.parse(new StringReader(synonymMappings));
+    return parser.build();
+  }
+}

src/test/java/com/yelp/nrtsearch/server/luceneserver/analysis/NrtsearchSynonymParserTest.java

@@ -0,0 +1,80 @@
+/*
+ * Copyright 2024 Yelp Inc.
+ *
+ * Licensed 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 com.yelp.nrtsearch.server.luceneserver.analysis;
+
+import static org.apache.lucene.analysis.BaseTokenStreamTestCase.assertAnalyzesTo;
+
+import com.carrotsearch.randomizedtesting.RandomizedRunner;
+import java.io.IOException;
+import java.io.StringReader;
+import java.text.ParseException;
+import org.apache.lucene.analysis.Analyzer;
+import org.apache.lucene.analysis.MockAnalyzer;
+import org.apache.lucene.analysis.MockTokenizer;
+import org.apache.lucene.analysis.Tokenizer;
+import org.apache.lucene.analysis.synonym.SynonymGraphFilter;
+import org.apache.lucene.analysis.synonym.SynonymMap;
+import org.apache.lucene.util.LuceneTestCase;
+import org.junit.runner.RunWith;
+
+@RunWith(RandomizedRunner.class)
+public class NrtsearchSynonymParserTest extends LuceneTestCase {
+
+  public void testParse() throws IOException, ParseException {
+    Analyzer analyzer = new MockAnalyzer(random());
+    NrtsearchSynonymParser parser =
+        new NrtsearchSynonymParser(Boolean.TRUE, Boolean.TRUE, analyzer);
+    String synonyms =
+        "a, b|ix, pie-ix|plaza, pla\\xE7a|plaza, plz|str, strada|str, strasse|str, stra\\xDFe|village, vlg";
+    parser.parse(new StringReader(synonyms));
+    final SynonymMap map = parser.build();
+    analyzer.close();
+
+    analyzer =
+        new Analyzer() {
+          @Override
+          protected TokenStreamComponents createComponents(String fieldName) {
+            Tokenizer tokenizer = new MockTokenizer(MockTokenizer.WHITESPACE, true);
+            return new TokenStreamComponents(
+                tokenizer, new SynonymGraphFilter(tokenizer, map, true));
+          }
+        };
+
+    assertAnalyzesTo(analyzer, "a", new String[] {"b", "a"}, new int[] {1, 0});
+    assertAnalyzesTo(
+        analyzer, "plaza", new String[] {"plaxe7a", "plz", "plaza"}, new int[] {1, 0, 0});
+    assertAnalyzesTo(
+        analyzer,
+        "str",
+        new String[] {"strada", "strasse", "straxdfe", "str"},
+        new int[] {1, 0, 0, 0});
+
+    analyzer.close();
+  }
+
+  public void testInvalidMappings() {
+    Analyzer analyzer = new MockAnalyzer(random());
+    NrtsearchSynonymParser parser =
+        new NrtsearchSynonymParser(Boolean.TRUE, Boolean.TRUE, analyzer);
+    String synonyms = "a, b, c, d, e";
+    expectThrows(
+        IllegalArgumentException.class,
+        () -> {
+          parser.parse(new StringReader(synonyms));
+        });
+    analyzer.close();
+  }
+}

src/test/java/com/yelp/nrtsearch/server/luceneserver/analysis/SynonymV2GraphFilterFactoryTest.java

@@ -0,0 +1,82 @@
+/*
+ * Copyright 2024 Yelp Inc.
+ *
+ * Licensed 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 com.yelp.nrtsearch.server.luceneserver.analysis;
+
+import static org.apache.lucene.util.LuceneTestCase.random;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
+import com.carrotsearch.randomizedtesting.RandomizedRunner;
+import java.io.IOException;
+import java.io.StringReader;
+import java.text.ParseException;
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.lucene.analysis.*;
+import org.apache.lucene.analysis.standard.StandardAnalyzer;
+import org.apache.lucene.analysis.tokenattributes.CharTermAttribute;
+import org.apache.lucene.util.LuceneTestCase;
+import org.junit.runner.RunWith;
+
+@RunWith(RandomizedRunner.class)
+public class SynonymV2GraphFilterFactoryTest extends LuceneTestCase {
+  public void testNoMappings() throws IOException, ParseException {
+    Analyzer analyzer = new MockAnalyzer(random());
+    try {
+      new SynonymV2GraphFilterFactory(new HashMap<>(), analyzer);
+      fail();
+    } catch (IllegalArgumentException e) {
+      assertEquals("Synonym mappings must be specified", e.getMessage());
+    }
+  }
+
+  public void testSingleMapping() throws IOException, ParseException {
+    SynonymV2GraphFilterFactory synonymV2GraphFilterFactory = getFactory("a, b");
+    TokenStream tokenStream =
+        new StandardAnalyzer().tokenStream("field", new StringReader("this is a test string"));
+    String[] expectedTokens = {"this", "is", "b", "a", "test", "string"};
+    assertTokenStream(synonymV2GraphFilterFactory, tokenStream, expectedTokens);
+  }
+
+  private static void assertTokenStream(
+      SynonymV2GraphFilterFactory synonymV2GraphFilterFactory,
+      TokenStream tokenStream,
+      String[] expectedTokens) {
+    try {
+      TokenStream output = synonymV2GraphFilterFactory.create(tokenStream);
+      CharTermAttribute charTermAtt = output.addAttribute(CharTermAttribute.class);
+      int i = 0;
+      output.reset();
+      while (output.incrementToken()) {
+        assertEquals(expectedTokens[i], charTermAtt.toString());
+        i += 1;
+      }
+      output.end();
+      output.close();
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  private SynonymV2GraphFilterFactory getFactory(String mappings)
+      throws IOException, ParseException {
+    Analyzer analyzer = new MockAnalyzer(random());
+    Map<String, String> params = new HashMap<>();
+    params.put(SynonymV2GraphFilterFactory.MAPPINGS, mappings);
+    params.put("parserFormat", "nrtsearch");
+    return new SynonymV2GraphFilterFactory(params, analyzer);
+  }
+}