Support collapse fields (PR 1/2) (#1276)

Author: papa99do

src/marqo/core/models/marqo_index.py

@@ -18,11 +18,21 @@
 # TODO refactor to remove dep to s2_inference
 from marqo.s2_inference import s2_inference
 from marqo.s2_inference.errors import UnknownModelError, InvalidModelPropertiesError
-import marqo.core.constants as constants
 
 logger = get_logger(__name__)
 
 
+class CollapseField(StrictBaseModel):
+    name: str
+    min_groups: int = pydantic.Field(default=500, gt=0, alias='minGroups')
+
+    @validator('name')
+    def validate_field_name_collapse(cls, v):
+        # Use common field name validation
+        validate_field_name(v)
+        return v
+
+
 class IndexType(Enum):
     Structured = 'structured'
     Unstructured = 'unstructured'
@@ -520,6 +530,7 @@ class SemiStructuredMarqoIndex(UnstructuredMarqoIndex):
     tensor_fields: List[TensorField]
     string_array_fields: Optional[List[
         StringArrayField]]  # This is required so that when saving a document containing string array fields, we can make changes to the schema on the fly. Ref: https://github.com/marqo-ai/marqo/blob/cfea70adea7039d1586c94e36adae8e66cabe306/src/marqo/core/semi_structured_vespa_index/semi_structured_vespa_schema_template_2_16.sd.jinja2#L83
+    collapse_fields: Optional[List[CollapseField]] = None
 
     def __init__(self, **data):
         super().__init__(**data)
@@ -528,6 +539,18 @@ def __init__(self, **data):
     def _valid_type(cls) -> IndexType:
         return IndexType.SemiStructured
 
+    @root_validator
+    def validate_collapse_fields(cls, values):
+        collapse_fields = values.get('collapse_fields')
+        if collapse_fields is not None and len(collapse_fields) != 1:
+            raise ValueError("There must be exactly one collapse field")
+        return values
+
+    def is_collapse_field(self, field_name: str) -> bool:
+        if not self.collapse_fields:
+            return False
+        return field_name in [field.name for field in self.collapse_fields]
+
     @property
     def field_map(self) -> Dict[str, Field]:
         """

src/marqo/core/models/marqo_index_request.py

@@ -41,6 +41,17 @@ class UnstructuredMarqoIndexRequest(MarqoIndexRequest):
     treat_urls_and_pointers_as_images: bool
     treat_urls_and_pointers_as_media: bool
     filter_string_max_length: int
+    collapse_fields: Optional[List[marqo_index.CollapseField]] = None
+
+    @root_validator
+    def validate_collapse_fields(cls, values):
+        collapse_fields = values.get('collapse_fields')
+        if collapse_fields is not None:
+            if len(collapse_fields) == 0:
+                raise ValueError("collapse_fields cannot be an empty list")
+            if len(collapse_fields) > 1:
+                raise ValueError("Only one collapse field is supported")
+        return values
 
 
 class FieldRequest(StrictBaseModel):

src/marqo/core/semi_structured_vespa_index/semi_structured_add_document_handler.py

@@ -43,6 +43,43 @@ def __init__(self, marqo_index: SemiStructuredMarqoIndex, add_docs_params: AddDo
         self.should_update_index = False
         self.field_count_config = field_count_config
 
+    def _validate_doc(self, doc) -> None:
+        """Override parent validation to add collapse field validation."""
+        # Call parent validation first
+        super()._validate_doc(doc)
+        
+        # Add collapse field validation
+        self._validate_collapse_field_presence(doc)
+
+    def _validate_collapse_field_presence(self, doc: Dict[str, Any]) -> None:
+        """Validate that documents contain required collapse fields with correct type."""
+        if not self.marqo_index.collapse_fields:
+            return  # No collapse fields configured
+        
+        collapse_field = self.marqo_index.collapse_fields[0]  # Only one allowed per spec
+        collapse_field_name = collapse_field.name
+
+        # TODO confirm if all these validations are required
+        if collapse_field_name not in doc:
+            raise AddDocumentsError(
+                f"Document missing required field '{collapse_field_name}'. "
+                f"All documents must contain this field for grouping."
+            )
+        
+        collapse_value = doc[collapse_field_name]
+        
+        if not isinstance(collapse_value, str):
+            raise AddDocumentsError(
+                f"Field '{collapse_field_name}' must be of type string. "
+                f"Got {type(collapse_value).__name__}: {collapse_value}"
+            )
+        
+        if not collapse_value.strip():
+            raise AddDocumentsError(
+                f"Field '{collapse_field_name}' cannot be empty. "
+                f"Provide a non-empty string value for grouping."
+            )
+
     def _handle_field(self, marqo_doc, field_name, field_content):
         """Handle a field in a Marqo document by processing it and updating the index schema if needed.
         
@@ -54,6 +91,10 @@ def _handle_field(self, marqo_doc, field_name, field_content):
         # Process field using parent class handler
         super()._handle_field(marqo_doc, field_name, field_content)
 
+        # Skip automatic lexical field creation for collapse fields - they are predefined in schema
+        if self.marqo_index.is_collapse_field(field_name):
+            return
+
         # Add lexical field if content is a string
         if isinstance(marqo_doc[field_name], str):
             language = self._get_field_language(field_name)

src/marqo/core/semi_structured_vespa_index/semi_structured_document.py

@@ -85,7 +85,9 @@ def from_vespa_document(cls, document: Dict, marqo_index: SemiStructuredMarqoInd
         string_array_prefix_length = len(common.STRING_ARRAY + '_')
 
         for field_name, field_value in fields.items():
-            if field_name in tensor_subfield_map:
+            if marqo_index.is_collapse_field(field_name):
+                text_fields[field_name] = field_value
+            elif field_name in tensor_subfield_map:
                 tensor_fields[field_name] = field_value
             elif field_name in lexical_field_map:
                 # Lexical fields are returned with prefixed name from get_by_ids
@@ -199,6 +201,10 @@ def _handle_field_content(cls, field_name: str, field_content: Union[str, bool,
 
     @classmethod
     def _handle_string_field(cls, field_name: str, field_content: str, instance, marqo_index: SemiStructuredMarqoIndex):
+        if marqo_index.is_collapse_field(field_name):
+            instance.text_fields[field_name] = field_content
+            return
+
         if field_name not in marqo_index.field_map:
             raise MarqoDocumentParsingError(f'Field {field_name} is not in index {marqo_index.name}')
 
@@ -209,6 +215,7 @@ def _handle_string_field(cls, field_name: str, field_content: str, instance, mar
             instance.fixed_fields.short_string_fields[field_name] = field_content
 
         if instance.index_supports_partial_updates:
+            # TODO do we need to store field type of collapse field? maybe since we need to support partial updates
             instance.fixed_fields.field_types[field_name] = MarqoFieldTypes.STRING.value
 
     @classmethod

src/marqo/core/semi_structured_vespa_index/semi_structured_vespa_schema.py

@@ -31,7 +31,15 @@ def generate_vespa_schema(cls, marqo_index: SemiStructuredMarqoIndex) -> str:
             vespa_schema_template = environment.get_template("semi_structured_vespa_schema_template_2_16.sd.jinja2")
         else:
             vespa_schema_template = environment.get_template("semi_structured_vespa_schema_template.sd.jinja2")
-        return vespa_schema_template.render(index=marqo_index, dimension=str(marqo_index.model.get_dimension()))
+
+        # simplify the logic in the template to just pass in the first collapse field if exists
+        collapse_field = marqo_index.collapse_fields[0] if marqo_index.collapse_fields else None
+
+        return vespa_schema_template.render(
+            index=marqo_index,
+            collapse_field=collapse_field,
+            dimension=str(marqo_index.model.get_dimension())
+        )
 
     def _generate_marqo_index(self, schema_name: str) -> SemiStructuredMarqoIndex:
         marqo_index = SemiStructuredMarqoIndex(
@@ -55,6 +63,7 @@ def _generate_marqo_index(self, schema_name: str) -> SemiStructuredMarqoIndex:
             filter_string_max_length=self._index_request.filter_string_max_length,
             treat_urls_and_pointers_as_images=self._index_request.treat_urls_and_pointers_as_images,
             treat_urls_and_pointers_as_media=self._index_request.treat_urls_and_pointers_as_media,
+            collapse_fields=self._index_request.collapse_fields,
         )
 
         return marqo_index

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

@@ -71,6 +71,14 @@ schema {{ index.schema_name }} {
             indexing: summary
         }
 
+        {% if collapse_field -%}
+        field {{ collapse_field.name }} type string {
+            indexing: attribute | summary
+            attribute: fast-search
+            rank: filter
+        }
+        {% endif -%}
+
         {% for lexical_field in index.lexical_fields -%}
         field {{ lexical_field.lexical_field_name }} type string {
             {%- if lexical_field.language %}
@@ -227,7 +235,19 @@ schema {{ index.schema_name }} {
             expression: sort_field_value(query(marqo__sort_field_weights_2))
         }
 
-        match-features: global_mult_modifier global_add_modifier sort_field_value_0  sort_field_value_1 sort_field_value_2
+        {% if collapse_field -%}
+        function collapse_field_hash() {
+            expression: attribute({{ collapse_field.name }})
+        }
+        {% endif -%}
+        match-features {
+            global_mult_modifier
+            global_add_modifier
+            sort_field_value_0
+            sort_field_value_1
+            sort_field_value_2
+            {% if collapse_field %}collapse_field_hash{% endif %}
+        }
     }
 
     {% if index.lexical_fields -%}
@@ -236,6 +256,17 @@ schema {{ index.schema_name }} {
             expression: modify(lexical_score(), query(marqo__mult_weights_lexical), query(marqo__add_weights_lexical))
         }
     }
+    {% if collapse_field -%}
+    rank-profile bm25_diversity inherits bm25 {
+        diversity {
+            attribute: {{ collapse_field.name }}
+            min-groups: {{ collapse_field.min_groups }}
+        }
+        second-phase {
+            expression: firstPhase
+        }
+    }
+    {% endif -%}
     {% endif -%}
 
     {# We provide this rank profile even without the tensor field to support embed requests -#}
@@ -255,6 +286,18 @@ schema {{ index.schema_name }} {
         {%- endif %}
     }
 
+    {% if collapse_field -%}
+    rank-profile embedding_similarity_diversity inherits embedding_similarity {
+        diversity {
+            attribute: {{ collapse_field.name }}
+            min-groups: {{ collapse_field.min_groups }}
+        }
+        second-phase {
+            expression: firstPhase
+        }
+    }
+    {% endif -%}
+
     {% if index.lexical_fields and index.tensor_fields -%}
     rank-profile hybrid_custom_searcher inherits default {
         inputs {
@@ -295,6 +338,26 @@ schema {{ index.schema_name }} {
             expression: modify(lexical_score(), query(marqo__mult_weights_lexical), query(marqo__add_weights_lexical))
         }
     }
+
+    {% if collapse_field -%}
+    rank-profile hybrid_bm25_then_embedding_similarity_diversity inherits hybrid_bm25_then_embedding_similarity {
+        diversity {
+            attribute: {{ collapse_field.name }}
+            min-groups: {{ collapse_field.min_groups }}
+        }
+    }
+
+    rank-profile hybrid_embedding_similarity_then_bm25_diversity inherits hybrid_embedding_similarity_then_bm25 {
+        diversity {
+            attribute: {{ collapse_field.name }}
+            min-groups: {{ collapse_field.min_groups }}
+        }
+        second-phase {
+            expression: firstPhase
+        }
+    }
+    {% endif -%}
+
     {%- endif %}
 
     document-summary all-non-vector-summary {
@@ -303,6 +366,9 @@ schema {{ index.schema_name }} {
         summary marqo__bool_fields type map<string, byte> {}
         summary marqo__int_fields type map<string, long> {}
         summary marqo__float_fields type map<string, double> {}
+        {% if collapse_field -%}
+        summary {{ collapse_field.name }} type string {}
+        {% endif -%}
         {% for string_array_field in index.string_array_fields -%}
         summary {{ string_array_field.string_array_field_name }} type array<string> {source: {{ string_array_field.string_array_field_name }}}
         {% endfor -%}
@@ -320,6 +386,9 @@ schema {{ index.schema_name }} {
         summary marqo__bool_fields type map<string, byte> {}
         summary marqo__int_fields type map<string, long> {}
         summary marqo__float_fields type map<string, double> {}
+        {% if collapse_field -%}
+        summary {{ collapse_field.name }} type string {}
+        {% endif -%}
         {% for string_array_field in index.string_array_fields -%}
         summary {{ string_array_field.string_array_field_name }} type array<string> {source: {{ string_array_field.string_array_field_name }}}
         {% endfor -%}

src/marqo/core/vespa_index/add_documents_handler.py

@@ -175,14 +175,17 @@ def add_documents(self) -> MarqoAddDocumentsResponse:
         with RequestMetricsStore.for_request().time("add_documents.vespa.to_vespa_docs"):
             vespa_docs = self._convert_to_vespa_docs()
 
-        self._pre_persist_to_vespa()
+        if vespa_docs:  # only continue if there's still vespa docs to persist
+            self._pre_persist_to_vespa()
 
-        # persist to vespa if there are still valid docs
-        with RequestMetricsStore.for_request().time("add_documents.vespa._bulk"):
-            response = self.vespa_client.feed_batch(vespa_docs, self.marqo_index.schema_name)
+            # persist to vespa if there are still valid docs
+            with RequestMetricsStore.for_request().time("add_documents.vespa._bulk"):
+                response = self.vespa_client.feed_batch(vespa_docs, self.marqo_index.schema_name)
 
-        with RequestMetricsStore.for_request().time("add_documents.postprocess"):
-            self._handle_vespa_response(response)
+            with RequestMetricsStore.for_request().time("add_documents.postprocess"):
+                self._handle_vespa_response(response)
+        else:
+            logger.debug('Skipping the Vespa roundtrip since there is no valid doc to feed')
 
         return self.add_docs_response_collector.to_add_doc_responses(self.marqo_index.name)
 

src/marqo/tensor_search/models/index_settings.py

@@ -23,6 +23,7 @@ class IndexSettings(StrictBaseModel):
     treatUrlsAndPointersAsImages: Optional[bool]
     treatUrlsAndPointersAsMedia: Optional[bool]
     filterStringMaxLength: Optional[int]
+    collapseFields: Optional[List[core.CollapseField]] = None
     model: str = 'hf/e5-base-v2'
     modelProperties: Optional[Dict[str, Any]]
     textQueryPrefix: Optional[str] = None
@@ -74,6 +75,19 @@ def validate_url_pointer_treatment(cls, values):
 
         return values
 
+    @root_validator
+    def validate_collapse_fields(cls, values):
+        collapse_fields = values.get('collapseFields')
+        index_type = values.get('type')
+        
+        # collapseFields is only supported for SemiStructuredIndex
+        if collapse_fields is not None and index_type == core.IndexType.Structured:
+            raise api_exceptions.InvalidArgError(
+                "collapseFields is only supported for unstructured indexes"
+            )
+        
+        return values
+
     @root_validator(pre=True)
     def validate_field_names(cls, values):
         # Verify no snake case field names (pydantic won't catch these due to allow_population_by_field_name = True)
@@ -191,6 +205,7 @@ def to_marqo_index_request(self, index_name: str) -> MarqoIndexRequest:
                 treat_urls_and_pointers_as_images=self.treatUrlsAndPointersAsImages,
                 treat_urls_and_pointers_as_media=self.treatUrlsAndPointersAsMedia,
                 filter_string_max_length=self.filterStringMaxLength,
+                collapse_fields=self.collapseFields,
                 marqo_version=version.get_version(),
                 created_at=time.time(),
                 updated_at=time.time()
@@ -204,11 +219,18 @@ def from_marqo_index(cls, marqo_index: core.MarqoIndex) -> "IndexSettings":
             # This covers both UnstructuredMarqoIndex and SemiStructuredMarqoIndex
             # We intentionally hide the lexical and tensor fields info in SemiStructuredMarqoIndex from customers since
             # this information and the SemiStructured concept are internal implementation details only.
+            
+            # Only include collapseFields for SemiStructuredMarqoIndex
+            collapse_fields = None
+            if isinstance(marqo_index, core.SemiStructuredMarqoIndex):
+                collapse_fields = marqo_index.collapse_fields
+            
             return cls(
                 type=core.IndexType.Unstructured,
                 treatUrlsAndPointersAsImages=marqo_index.treat_urls_and_pointers_as_images,
                 treatUrlsAndPointersAsMedia=marqo_index.treat_urls_and_pointers_as_media,
                 filterStringMaxLength=marqo_index.filter_string_max_length,
+                collapseFields=collapse_fields,
                 model=marqo_index.model.name,
                 modelProperties=IndexSettings.get_model_properties(marqo_index),
                 normalizeEmbeddings=marqo_index.normalize_embeddings,