Merge pull request #861 from hubmapconsortium/nickakhmetov/immediate-…

…relation-ids

Nickakhmetov/immediate relation ids

src/hubmap_translation/addl_index_transformations/portal/translate.py

@@ -23,6 +23,7 @@ def translate(doc, organ_map):
     _translate_timestamp(doc)
     _translate_access_level(doc)
     _translate_external_consortium(doc)
+    _add_spatial_info(doc)
 
 
 def _map(doc, key, map):
@@ -301,3 +302,31 @@ def _add_origin_samples_unique_mapped_organs(doc):
     if doc['entity_type'] in ['Sample', 'Dataset'] and 'origin_samples' in doc:
         doc['origin_samples_unique_mapped_organs'] = _get_unique_mapped_organs(
             doc['origin_samples'])
+
+
+def _add_spatial_info(doc):
+    '''
+    Add a boolean field "is_spatial" to the document based on the entity type and the presence of an rui_location field.
+
+    For samples, the is_spatial field is set to True if the rui_location field is present.
+    >>> doc = {'entity_type': 'Sample', 'rui_location': 'https://example.com'}
+    >>> _add_spatial_info(doc); doc['is_spatial']
+    True
+
+    For datasets, the is_spatial field is set to True if any ancestor has an rui_location field.
+    The rui_location field is also copied from the nearest ancestor with an rui_location field.
+    >>> doc = {'entity_type': 'Dataset', 'ancestors': [{'rui_location': 'https://example.com'}, {'rui_location': 'https://example2.com'}]}
+    >>> _add_spatial_info(doc); doc['is_spatial']; doc['rui_location']
+    True
+    'https://example2.com'
+    '''
+    if (doc['entity_type'] == 'Sample'):
+        doc['is_spatial'] = doc.get('rui_location', None) is not None
+    if (doc['entity_type'] == 'Dataset'):
+        ancestors = doc.get('ancestors', [])
+        # Find the nearest ancestor with an rui_location - the last one in the list with an rui_location field.
+        nearest_rui_location_ancestor = next(
+            (ancestor for ancestor in reversed(ancestors) if 'rui_location' in ancestor), None)
+        if nearest_rui_location_ancestor is not None:
+            doc['is_spatial'] = nearest_rui_location_ancestor is not None
+            doc['rui_location'] = nearest_rui_location_ancestor['rui_location']

src/hubmap_translator.py

@@ -39,6 +39,8 @@
     'datasets',
     'immediate_ancestors',
     'immediate_descendants'
+    'immediate_ancestor_ids',
+    'immediate_descendant_ids'
 ]
 
 # A map keyed by entity attribute names stored in Neo4j and retrieved from entity-api, with
@@ -1457,6 +1459,9 @@ def generate_doc(self, entity, return_type):
                 entity['immediate_ancestors'] = immediate_ancestors
                 entity['immediate_descendants'] = immediate_descendants
 
+                entity['immediate_ancestor_ids'] = immediate_ancestor_ids
+                entity['immediate_descendant_ids'] = immediate_descendant_ids
+
             # The `sample_category` is "organ" and the `organ` code is set at the same time
             if entity['entity_type'] in ['Sample', 'Dataset', 'Publication']:
                 # Add new properties