Implement configurable stemming (#1273)

Implement configurable text field stemming.

This implementation is similar to language. Stemming can be configured per text field (and only for text fields). It's configured in the Vespa schema. Stemming cannot be changed once the field has been created. Unlike language, search API does not support changing stemming. It's picked by Vespa based on the fields being searched.

Note only full text search is stemmed. Filtering isn't affected.

Author: farshidz

scripts/vespa_local/vespa_local.py

@@ -157,6 +157,7 @@ def get_services_xml_content(self) -> str:
                     <search/>
                     <nodes>
                         <node hostalias="node1"/>
+                        <jvm options="-Xms32M -Xmx256M -agentlib:jdwp=transport=dt_socket,server=y,suspend=n,address=*:5005" />
                     </nodes>
                 </container>
                 <content id="content_default" version="1.0">

src/marqo/core/constants.py

@@ -26,6 +26,7 @@
 MARQO_RERANK_DEPTH_MINIMUM_VERSION = semver.VersionInfo.parse('2.15.0')
 MARQO_SORT_BY_MINIMUM_VERSION = semver.VersionInfo.parse('2.22.0')
 MARQO_LANGUAGE_MINIMUM_VERSION = semver.VersionInfo.parse('2.16.0')
+MARQO_STEMMING_MINIMUM_VERSION = semver.VersionInfo.parse('2.16.0')
 MARQO_PARTIAL_UPDATE_MINIMUM_VERSION = semver.VersionInfo.parse('2.16.0')
 
 # For score modifiers
@@ -37,4 +38,4 @@
 QUERY_INPUT_SCORE_MODIFIERS_ADD_WEIGHTS_TENSOR = 'marqo__add_weights_tensor'
 QUERY_INPUT_SCORE_MODIFIERS_MULT_WEIGHTS_GLOBAL = 'marqo__mult_weights_global'
 QUERY_INPUT_SCORE_MODIFIERS_ADD_WEIGHTS_GLOBAL = 'marqo__add_weights_global'
-MARQO_GLOBAL_SCORE_MODIFIERS = 'global'
+MARQO_GLOBAL_SCORE_MODIFIERS = 'global'

src/marqo/core/index_management/index_management.py

@@ -219,24 +219,25 @@ def update_index(self, marqo_index: SemiStructuredMarqoIndex) -> None:
             OperationConflictError: If another index creation/deletion operation is
                 in progress and the lock cannot be acquired
         """
-        with self._vespa_deployment_lock():
-            existing_index = self.get_index(marqo_index.name)
-            if not isinstance(existing_index, SemiStructuredMarqoIndex):
-                # This is just a sanity check, it should not happen since we do not expose this method to end user.
-                raise InternalError(f'Index {marqo_index.name} created by Marqo version {marqo_index.marqo_version} '
-                                    f'can not be updated.')
-
-            def is_subset(dict_a, dict_b):
-                # check if dict_a is a subset of dict_b
-                return all(k in dict_b and dict_b[k] == v for k, v in dict_a.items())
-
-            if (is_subset(marqo_index.tensor_field_map, existing_index.tensor_field_map) and
-                    is_subset(marqo_index.field_map, existing_index.field_map) and
-                        is_subset(marqo_index.name_to_string_array_field_map, existing_index.name_to_string_array_field_map)):
-                logger.debug(f'Another thread has updated the index {marqo_index.name} already.')
-                return
+        existing_index = self.get_index(marqo_index.name)
+        if not isinstance(existing_index, SemiStructuredMarqoIndex):
+            # This is just a sanity check, it should not happen since we do not expose this method to end user.
+            raise InternalError(f'Index {marqo_index.name} created by Marqo version {marqo_index.marqo_version} '
+                                f'can not be updated.')
+
+        def is_subset(dict_a, dict_b):
+            # check if dict_a is a subset of dict_b
+            return all(k in dict_b and dict_b[k] == v for k, v in dict_a.items())
+
+        if (is_subset(marqo_index.tensor_field_map, existing_index.tensor_field_map) and
+                is_subset(marqo_index.field_map, existing_index.field_map) and
+                    is_subset(marqo_index.name_to_string_array_field_map, existing_index.name_to_string_array_field_map)):
+            logger.debug(f'Another thread has updated the index {marqo_index.name} already.')
+            return
 
+        with self._vespa_deployment_lock():
             schema = SemiStructuredVespaSchema.generate_vespa_schema(marqo_index)
+            logger.debug(f'Updating index {marqo_index.name} with schema:\n{schema}')
             self._get_vespa_application().update_index_setting_and_schema(marqo_index, schema)
 
     def _get_existing_indexes(self) -> List[MarqoIndex]:

src/marqo/core/models/marqo_index.py

@@ -73,6 +73,13 @@ class FieldFeature(Enum):
     Filter = 'filter'
 
 
+class Stemming(str, Enum):
+    None_ = 'none'
+    Best = 'best'
+    Shortest = 'shortest'
+    Multiple = 'multiple'
+
+
 class DistanceMetric(Enum):
     Euclidean = 'euclidean'
     Angular = 'angular'
@@ -105,6 +112,7 @@ class Field(ImmutableStrictBaseModel):
     filter_field_name: Optional[str]
     dependent_fields: Optional[Dict[str, float]]
     language: Optional[str] = None
+    stemming: Optional[Stemming] = None
 
     @root_validator
     def check_all_fields(cls, values):
@@ -660,6 +668,15 @@ def index_supports_language(self) -> bool:
             'index_supports_language',
             lambda: self.parsed_marqo_version() >= constants.MARQO_LANGUAGE_MINIMUM_VERSION)
 
+    @property
+    def index_supports_stemming(self) -> bool:
+        """
+        Check if the index supports stemming.
+        """
+        return self._cache_or_get(
+            'index_supports_stemming',
+            lambda: self.parsed_marqo_version() >= constants.MARQO_STEMMING_MINIMUM_VERSION)
+
     @property
     def index_supports_sorty_by(self) -> bool:
         """
@@ -733,6 +750,7 @@ def validate_structured_field(values, marqo_index: bool) -> None:
     type: FieldType = values['type']
     features: List[FieldFeature] = values['features']
     language: str = values.get('language')
+    stemming: str = values.get('stemming')
     dependent_fields: Optional[Dict[str, float]] = values['dependent_fields']
 
     validate_field_name(name)
@@ -762,6 +780,12 @@ def validate_structured_field(values, marqo_index: bool) -> None:
             f'feature is present'
         )
 
+    if stemming is not None and FieldFeature.LexicalSearch not in features:
+        raise ValueError(
+            f'{name}: stemming can only be populated when {FieldFeature.LexicalSearch.value} '
+            f'feature is present'
+        )
+
     if FieldFeature.ScoreModifier in features and type not in [FieldType.Float, FieldType.Int,
                                                                FieldType.Double, FieldType.MapFloat,
                                                                FieldType.MapInt, FieldType.MapDouble,

src/marqo/core/semi_structured_vespa_index/semi_structured_add_document_handler.py

@@ -1,4 +1,4 @@
-from typing import Dict, Any
+from typing import Dict, Any, Optional
 
 import pydantic.v1 as pydantic
 
@@ -98,7 +98,8 @@ def _handle_field(self, marqo_doc, field_name, field_content):
         # Add lexical field if content is a string
         if isinstance(marqo_doc[field_name], str):
             language = self._get_field_language(field_name)
-            self._add_lexical_field_to_index(field_name, language)
+            stemming = self._get_field_stemming(field_name)
+            self._add_lexical_field_to_index(field_name, language, stemming)
 
         # Add string array field if content is list of strings and index version supports it
         is_string_array = (
@@ -145,27 +146,53 @@ def _get_field_language(self, field_name):
             return None
 
         if field_mapping.get('type') == 'text_field':
-            if not self.marqo_index.index_supports_language:
+            language = field_mapping.get('language')
+            if language is not None and not self.marqo_index.index_supports_language:
                 raise AddDocumentsError(
                     f'Language is only supported for indexes created with Marqo version '
                     f'{constants.MARQO_LANGUAGE_MINIMUM_VERSION} or later. This index was created with  '
                     f'Marqo {self.marqo_index.marqo_version}.'
                 )
-            return field_mapping.get('language')
+            return language
 
         return None
 
+    def _get_field_stemming(self, field_name) -> Optional[str]:
+        """Extract stemming specification for a field from mappings and validate."""
+        if not self.add_docs_params.mappings:
+            return None
 
-    def _add_lexical_field_to_index(self, field_name, language=None):
+        field_mapping = self.add_docs_params.mappings.get(field_name)
+        if not field_mapping:
+            return None
+
+        if field_mapping.get('type') == 'text_field':
+            stemming = field_mapping.get('stemming')
+            if stemming is not None and not self.marqo_index.index_supports_stemming:
+                raise AddDocumentsError(
+                    f'Stemming is only supported for indexes created with Marqo version '
+                    f'{constants.MARQO_STEMMING_MINIMUM_VERSION} or later. This index was created with  '
+                    f'Marqo {self.marqo_index.marqo_version}.'
+                )
+            return stemming
+
+        return None
+
+    def _add_lexical_field_to_index(self, field_name, language=None, stemming=None):
         if field_name in self.marqo_index.field_map:
-            if language is not None:
-                existing_field = self.marqo_index.field_map[field_name]
-                if existing_field.language != language:
-                    raise AddDocumentsError(
-                        f"Field '{field_name}' already exists with a different language configuration. "
-                        f"Cannot change language from '{existing_field.language}' to '{language}' "
-                        f"for existing field."
-                    )
+            existing_field = self.marqo_index.field_map[field_name]
+            if language is not None and existing_field.language != language:
+                raise AddDocumentsError(
+                    f"Field '{field_name}' already exists with a different language configuration. "
+                    f"Cannot change language from '{existing_field.language}' to '{language}' "
+                    f"for existing field."
+                )
+            if stemming is not None and existing_field.stemming != stemming:
+                raise AddDocumentsError(
+                    f"Field '{field_name}' already exists with a different stemming configuration. "
+                    f"Cannot change stemming from '{existing_field.stemming}' to '{stemming}' "
+                    f"for existing field."
+                )
             return
 
         max_lexical_field_count = self.field_count_config.max_lexical_field_count
@@ -176,14 +203,18 @@ def _add_lexical_field_to_index(self, field_name, language=None):
                                      f'limit in MARQO_MAX_LEXICAL_FIELD_COUNT_UNSTRUCTURED environment variable.')
 
         # Add missing lexical fields to marqo index
-        logger.debug(f'Adding lexical field {field_name} to index {self.marqo_index.name}' +
-                    (f' with language {language}' if language else ''))
+        debug_parts = [f'Adding lexical field {field_name} to index {self.marqo_index.name}']
+        if language:
+            debug_parts.append(f'with language {language}')
+        if stemming:
+            debug_parts.append(f'with stemming {stemming}')
+        logger.debug(' '.join(debug_parts))
 
         self.marqo_index.lexical_fields.append(
             Field(name=field_name, type=FieldType.Text,
                   features=[FieldFeature.LexicalSearch],
                   lexical_field_name=f'{SemiStructuredVespaSchema.FIELD_INDEX_PREFIX}{field_name}',
-                  language=language)
+                  language=language, stemming=stemming)
         )
         self.marqo_index.clear_cache()
         self.should_update_index = True

src/marqo/core/semi_structured_vespa_index/semi_structured_vespa_schema_template_2_16.sd.jinja2

@@ -90,6 +90,9 @@ schema {{ index.schema_name }} {
             indexing: index | summary
             {%- endif %}
             index: enable-bm25
+            {%- if lexical_field.stemming %}
+            stemming: {{ lexical_field.stemming.value }}
+            {%- endif %}
         }
         {% endfor -%}
 

src/marqo/tensor_search/models/mappings_object.py

@@ -70,8 +70,16 @@
         "language": {
             "type": "string",
             "minLength": 1
+        },
+        "stemming": {
+            "type": "string",
+            "enum": ["none", "best", "shortest", "multiple"]
         }
     },
-    "required": ["type", "language"],
+    "required": ["type"],
+    "anyOf": [
+        {"required": ["language"]},
+        {"required": ["stemming"]}
+    ],
     "additionalProperties": False
 }

tests/api_tests/v1/tests/api_tests/test_language.py

@@ -33,19 +33,22 @@ def populate_index(self):
                 "_id": "1",
                 "title1": "Vestido Mole Perfeito",  # Portuguese
                 "title2": "Collections de livres",  # French
-                "title3": "White dog"  # English
+                "title3": "White dog",  # English,
+                "size": "M"
             },
             {
                 "_id": "2",
                 "title1": "Vestido Mole Confortável",
                 "title2": "collections art francais",
-                "title3": "black cat"
+                "title3": "black cat",
+                "size": "M"
             },
             {
                 "_id": "3",
                 "title1": "Vestido Leve e Elegante",
                 "title2": "mes collections Preferees",
-                "title3": "blue sky"
+                "title3": "blue sky",
+                "size": "S"
             }
         ]
 
@@ -175,6 +178,43 @@ def test_field_language_override(self):
         self.assertEqual(["1"], hits_pt, "Should find only the Portuguese doc")
         self.assertEqual(["2"], hits_en, "Should find only the English doc")
 
+    def test_facets_and_relevance_cutoff(self):
+        self.populate_index()
+
+        cases = [
+            ("pt", True),
+            ("en", False)
+        ]
+
+        for language, matches in cases:
+            with self.subTest(f"Testing facets and relevance cutoff for language: {language}"):
+                res = self.client.index(self.multilingual_index_name).search(
+                    q="mole",
+                    search_method="HYBRID",
+                    language=language,
+                    hybrid_parameters={
+                        'searchableAttributesLexical': ['title1'],
+                        'retrievalMethod': 'lexical',
+                        'rankingMethod': 'lexical',
+                    }, # Use lexical retrieval so that facets don't get tensor hits
+                    facets={
+                        "fields": {
+                            "size": {"type": "string"}
+                        }
+                    }, relevance_cutoff={"method": "mean_std_dev", "parameters": {"stdDevFactor": 1.2}},
+                )
+
+                if matches:
+                    self.assertGreater(len(res["hits"]), 0, "Should find matches for 'mole'")
+                    self.assertGreater(len(res['facets']['size']), 0, "Should have facets for 'size'")
+                    self.assertGreater(res["_relevantCandidates"], 0,
+                                       "Should have relevant candidates count greater than 0")
+                else:
+                    self.assertEqual(len(res["hits"]), 0, "Should find no matches for 'mole' in English")
+                    self.assertEqual(len(res['facets']), 0, "Should have no facets for 'size' in English")
+                    self.assertEqual(res["_relevantCandidates"], 0,
+                                     "Should have no relevant candidates count in English")
+
     def test_tensor_search_with_language_error(self):
         """Test that specifying language for tensor search raises an error."""
         self.populate_index()