Support collapse fields (PR 2/2) (#1277)

Author: papa99do

src/marqo/core/constants.py

@@ -28,6 +28,7 @@
 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')
+MARQO_COLLAPSE_FIELDS_MINIMUM_VERSION = semver.VersionInfo.parse('2.23.0')
 
 # For score modifiers
 QUERY_INPUT_SCORE_MODIFIERS_MULT_WEIGHTS_2_9 = 'marqo__mult_weights'

src/marqo/core/models/marqo_query.py

@@ -73,6 +73,7 @@ class MarqoHybridQuery(MarqoTensorQuery, MarqoLexicalQuery):
     track_total_hits: Optional[bool] = None
     relevance_cutoff: Optional[RelevanceCutoffModel] = None
     sort_by: Optional[SortByModel] = None
+    collapse_field_name: Optional[str] = None
 
     @root_validator(pre=True)
     def validate_searchable_attributes_and_score_modifiers(cls, values):

src/marqo/core/search/hybrid_search.py

@@ -50,7 +50,8 @@ def search(
             language: Optional[str] = None,
             relevance_cutoff: Optional[RelevanceCutoffModel] = None,
             sort_by: Optional[SortByModel] = None,
-            interpolation_method: Optional[InterpolationMethod] = None
+            interpolation_method: Optional[InterpolationMethod] = None,
+            collapse_field_name: Optional[str] = None
     ) -> Dict:
         """
 
@@ -81,6 +82,7 @@ def search(
                 relevance_cutoff: RelevanceCutoffModel object to specify relevance cutoff for the search.
                 sort_by: SortByModel object to specify sorting for the search. If not provided, no sorting will be applied.
                 interpolation_method: InterpolationMethod object to specify the interpolation method for hybrid search.
+                collapse_field_name:  field name to collapse the search result on.
             Returns:
 
             Output format:
@@ -292,7 +294,8 @@ def search(
             track_total_hits=track_total_hits,
             language=language,
             relevance_cutoff=relevance_cutoff,
-            sort_by=sort_by
+            sort_by=sort_by,
+            collapse_field_name=collapse_field_name
         )
 
         vespa_index = vespa_index_factory(marqo_index)

src/marqo/core/semi_structured_vespa_index/semi_structured_vespa_index.py

@@ -26,6 +26,8 @@
 class SemiStructuredVespaIndex(StructuredVespaIndex, UnstructuredVespaIndex):
     """
     An implementation of VespaIndex for SemiStructured indexes.
+    TODO the multi-inheritance makes the implementation difficult to reason about. Consider refactor to composition
+      instead. e.g. extract different logics to different query component builders, and combined the result.
     """
     index_supports_partial_updates: bool = False
 
@@ -108,8 +110,24 @@ def _to_vespa_hybrid_query(self, marqo_query):
         # add sort by and relevance cutoff
         self._add_relevance_cutoff_and_sort_by_params(marqo_query, query)
 
+        # add the collapse_field to query
+        if marqo_query.collapse_field_name:
+            query.update(self._generate_collapse_query_params(marqo_query.collapse_field_name))
+
         return query
 
+    def _generate_collapse_query_params(self, collapse_field_name: str):
+        return {
+            'collapsefield': collapse_field_name,
+            'collapsesize': 1,  # currently fixed to 1, will support multiple if needed in the future
+
+            # use a different rank profile to ensure diversity in the result returned to Vespa container
+            'marqo__ranking.lexical.lexical': common.RANK_PROFILE_BM25 + '_diversity',
+            'marqo__ranking.tensor.tensor': common.RANK_PROFILE_EMBEDDING_SIMILARITY + '_diversity',
+            'marqo__ranking.lexical.tensor': common.RANK_PROFILE_HYBRID_BM25_THEN_EMBEDDING_SIMILARITY + '_diversity',
+            'marqo__ranking.tensor.lexical': common.RANK_PROFILE_HYBRID_EMBEDDING_SIMILARITY_THEN_BM25 + '_diversity',
+        }
+
     def _add_relevance_cutoff_and_sort_by_params(self, marqo_query, query):
         if marqo_query.relevance_cutoff:
             query["marqo__hybrid.relevanceCutoff.method"] = marqo_query.relevance_cutoff.method
@@ -136,6 +154,8 @@ def _add_relevance_cutoff_and_sort_by_params(self, marqo_query, query):
             for index, field in enumerate(marqo_query.sort_by.fields):
                 query["query_features"][f'marqo__sort_field_weights_{index}'] = {field.field_name: 1}
 
+        return query
+
     def _generate_facet_queries(self, marqo_query):
         facets_query_skeleton = '%s limit 0 | %s'
         QUERY_DELIMITER = "\n---MARQO-YQL-QUERY-DELIMITER---\n"
@@ -171,7 +191,8 @@ def _generate_facet_queries(self, marqo_query):
             f'{base_yql}{filter_term}', f"all(group({self._TOTAL_HITS_GROUP_CONST}) each(output(count())))"))
 
         if marqo_query.facets is not None:
-            facets_term = self._get_facets_term(marqo_query.facets)
+            facets_term = self._get_facets_term(marqo_query.facets,
+                                                collapse_field_name=marqo_query.collapse_field_name)
 
             if facets_term is not None:
                 facet_queries.append(facets_query_skeleton % (f'{base_yql}{filter_term}', facets_term))
@@ -190,15 +211,17 @@ def _generate_facet_queries(self, marqo_query):
                         new_filter_term = f' AND {new_filter_term}'
                     else:
                         new_filter_term = ''
-                    new_facets_term = self._get_facets_term(marqo_query.facets, facet_parameters.exclude_terms)
+                    new_facets_term = self._get_facets_term(marqo_query.facets, facet_parameters.exclude_terms,
+                                                            collapse_field_name=marqo_query.collapse_field_name)
 
                     query_yql = f'{base_yql}{new_filter_term}'
 
                     facet_queries.append(facets_query_skeleton % (query_yql, new_facets_term))
 
         return QUERY_DELIMITER.join(facet_queries)
 
-    def _get_facets_term(self, facets_parameters: FacetsParameters, exclusion_terms: List[str] = None) -> str:
+    def _get_facets_term(self, facets_parameters: FacetsParameters, exclusion_terms: List[str] = None,
+                         collapse_field_name: Optional[str] = None) -> str:
         """
         Build a facets grouping query string from the provided facets_parameters.
         """
@@ -260,7 +283,11 @@ def build_field_group(field_config, field_name, field_id, field_type_overwrite=N
             params = build_group_parameters(field_config)
 
             # Build output expression
-            if field_config.type == "number":
+            if collapse_field_name:
+                # we can only get the count collapsed to this field regardless of data type
+                output = f"each(group({collapse_field_name}) output(count()))"
+            elif field_config.type == "number":
+                # if we do not collapse, we can get the following stats with count for number type
                 aggregations = ["sum", "avg", "min", "max"]
                 field_type = FIELD_TYPES[field_type_overwrite]
                 funcs = [f'{func}({field_type}{{"{field_name}"}})' for func in aggregations]
@@ -321,6 +348,10 @@ def generate_equality_filter_string(node: search_filter.EqualityTerm) -> str:
             if node.field == MARQO_DOC_ID:
                 return f'({VESPA_FIELD_ID} contains "{node.value}")'
 
+            if self.get_marqo_index().is_collapse_field(node.field):
+                # collapse field is indexed as attribute, can be used directly in a filter term
+                return f'({node.field} contains "{node.value}")'
+
             # Bool Filter
             if node.value.lower() in self._FILTER_STRING_BOOL_VALUES:
                 filter_value = int(True if node.value.lower() == "true" else False)
@@ -1034,11 +1065,15 @@ def _combine_number_stats(self, current_stats, stats):
         if current_stats == {}:
             return stats
         aggregated_stats = {}
-        aggregated_stats["sum"] = current_stats["sum"] + stats["sum"]
         aggregated_stats["count"] = current_stats["count"] + stats["count"]
-        aggregated_stats["avg"] = (current_stats["avg"] * current_stats["count"] + stats["avg"] * stats["count"]) / (current_stats["count"] + stats["count"])
-        aggregated_stats["min"] = min(current_stats["min"], stats["min"])
-        aggregated_stats["max"] = max(current_stats["max"], stats["max"])
+        if "sum" in current_stats and "sum" in stats:
+            aggregated_stats["sum"] = current_stats["sum"] + stats["sum"]
+        if "avg" in current_stats and "avg" in stats:
+            aggregated_stats["avg"] = (current_stats["avg"] * current_stats["count"] + stats["avg"] * stats["count"]) / (current_stats["count"] + stats["count"])
+        if "min" in current_stats and "min" in stats:
+            aggregated_stats["min"] = min(current_stats["min"], stats["min"])
+        if "max" in current_stats and "max" in stats:
+            aggregated_stats["max"] = max(current_stats["max"], stats["max"])
         return aggregated_stats
 
 

src/marqo/tensor_search/api.py

@@ -439,7 +439,8 @@ def search(index_name: str, search_query_dict: dict, device: str = Depends(api_v
                 language=search_query.language,
                 relevance_cutoff= search_query.relevance_cutoff,
                 sort_by = search_query.sort_by,
-                interpolation_method=search_query.interpolationMethod
+                interpolation_method=search_query.interpolationMethod,
+                collapse_field_name=search_query.collapse_fields[0].name if search_query.collapse_fields else None
             )
             return ORJSONResponse(result)
         except Exception as e:

src/marqo/tensor_search/models/api_models.py

@@ -29,6 +29,10 @@ class Config:
     pass
 
 
+class SearchCollapseField(ImmutableStrictBaseModel):
+    name: str
+
+
 class CustomVectorQuery(ImmutableStrictBaseModel):
     class CustomVector(ImmutableStrictBaseModel):
         content: Optional[str] = None
@@ -68,6 +72,7 @@ class Config(BaseMarqoModel.Config):
     sort_by: Optional[SortByModel] = Field(default=None, alias="sortBy")
     relevance_cutoff: Optional[RelevanceCutoffModel] = Field(default=None, alias="relevanceCutoff")
     interpolationMethod: Optional[InterpolationMethod] = None
+    collapse_fields: Optional[List[SearchCollapseField]] = Field(default=None, alias="collapseFields")
 
     # By default, we retrieve 3 times more candidates than the limit to ensure we have enough results to sort.
     _DEFAULT_SORT_CANDIDATES_MULTIPLIER = 3
@@ -424,6 +429,25 @@ def _validate_and_set_sort_by_min_sort_candidates_parameters(cls, values):
                 )
         return values
 
+    @root_validator(pre=False)
+    def validate_collapse_fields_only_for_hybrid_search(cls, values):
+        """Validate collapse fields only provided for hybrid search"""
+        collapse_fields = values.get('collapse_fields')
+        search_method = values.get('searchMethod')
+        if collapse_fields is not None and search_method.upper() != SearchMethod.HYBRID:
+            raise ValueError(f"collapseFields can only be provided for 'HYBRID' search. "
+                             f"Search method is {search_method}.")
+        return values
+        
+    @root_validator(pre=False) 
+    def validate_single_collapse_field(cls, values):
+        """Validate exactly one collapse field is provided"""
+        collapse_fields = values.get('collapse_fields')
+        if collapse_fields is not None:
+            if len(collapse_fields) != 1:
+                raise ValueError("Exactly one collapse field must be provided")
+        return values
+
 
 class BulkSearchQueryEntity(SearchQuery):
     index: MarqoIndex

src/marqo/tensor_search/tensor_search.py

@@ -359,7 +359,8 @@ def search(config: Config, index_name: str, text: Optional[Union[str, dict, Cust
            language: Optional[str] = None,
            relevance_cutoff: Optional[RelevanceCutoffModel] = None,
            sort_by: Optional[SortByModel] = None,
-           interpolation_method: Optional[InterpolationMethod] = None
+           interpolation_method: Optional[InterpolationMethod] = None,
+           collapse_field_name: Optional[str] = None
            ) -> Dict:
     """The root search method. Calls the specific search method
 
@@ -454,13 +455,32 @@ def search(config: Config, index_name: str, text: Optional[Union[str, dict, Cust
     # Fetch marqo index to pass to search method
     marqo_index = index_meta_cache.get_index(index_management=config.index_management, index_name=index_name)
     marqo_index_version = marqo_index.parsed_marqo_version()
+    
+    # Validate collapse field configuration
+    if collapse_field_name is not None:
+        # Validate if the index version support this feature
+        if (marqo_index_version < constants.MARQO_COLLAPSE_FIELDS_MINIMUM_VERSION or
+                not isinstance(marqo_index, SemiStructuredMarqoIndex)):
+            index_type = 'structured' if marqo_index.type == IndexType.Structured else 'unstructured'
+            raise core_exceptions.UnsupportedFeatureError(
+                f"The 'collapseFields' search parameter is only supported for unstructured indexes created with "
+                f"Marqo version {str(constants.MARQO_COLLAPSE_FIELDS_MINIMUM_VERSION)} or later. "
+                f"This index is {index_type} and was created with Marqo {marqo_index_version}."
+            )
+
+        # Validate collapse field exists in index configuration
+        if not marqo_index.is_collapse_field(collapse_field_name):
+            raise api_exceptions.InvalidArgError(f"Field '{collapse_field_name}' is not configured as a collapse field "
+                                                 f"for this index")
+    
     if rerank_depth is not None \
             and marqo_index_version < constants.MARQO_RERANK_DEPTH_MINIMUM_VERSION:
         raise core_exceptions.UnsupportedFeatureError(
             f"The 'rerankDepth' search parameter is only supported for indexes created with Marqo version "
             f"{str(constants.MARQO_RERANK_DEPTH_MINIMUM_VERSION)} or later. "
             f"This index was created with Marqo {marqo_index_version}."
         )
+
     if search_method.upper() in {SearchMethod.TENSOR, SearchMethod.HYBRID}:
         # Default approximate and efSearch -- we can't set these at API-level since they're not a valid args
         # for lexical search
@@ -511,7 +531,8 @@ def search(config: Config, index_name: str, text: Optional[Union[str, dict, Cust
                 hybrid_parameters=hybrid_parameters, facets=facets, track_total_hits=track_total_hits,
                 language=language,
                 relevance_cutoff=relevance_cutoff, sort_by=sort_by,
-                interpolation_method=interpolation_method
+                interpolation_method=interpolation_method,
+                collapse_field_name=collapse_field_name
             )
 
     elif search_method.upper() == SearchMethod.LEXICAL:

src/marqo/version.py

@@ -1,4 +1,4 @@
-__version__ = "2.22.1"
+__version__ = "2.23.0"
 
 def get_version() -> str:
     return f"{__version__}"

src/marqo/vespa/vespa_client.py

@@ -257,7 +257,8 @@ def query(self, yql: str, hits: int = 10, ranking: str = None, model_restrict: s
 
         self._query_raise_for_status(resp)
 
-        return QueryResult(**orjson.loads(resp.text))
+        resp_dict = orjson.loads(resp.text)
+        return QueryResult(**resp_dict)
 
     def feed_document(self, document: VespaDocument, schema: str, timeout: int = 60) -> FeedDocumentResponse:
         """