[FIX] flatten display of features and fix query list of lists (#881)

* flatten display of features and fix query list of lists

* fix tests

Author: jdkent

store/neurostore/resources/pipeline.py

@@ -52,7 +52,6 @@ def build_jsonpath(field_path: str, operator: str, value: str) -> str:
     Returns:
         PostgreSQL jsonpath query string
     """
-
     # Handle regular field queries
     cast_val, is_numeric = determine_value_type(value)
 
@@ -83,35 +82,35 @@ def build_jsonpath(field_path: str, operator: str, value: str) -> str:
             raw_value = f'"{cast_val}"'
         raw_value = f"@ {sql_op} {raw_value}"
 
-    # Check if we're querying an array field
+    # Check if we're querying array fields
     path_parts = field_path.split(".")
     if any(p.endswith("[]") for p in path_parts):
-        # Handle array field queries
+        query = "$"
         path_segments = []
-        for i, part in enumerate(path_parts):
+
+        for part in path_parts:
             if part.endswith("[]"):
-                # Convert path up to this point into the base path
-                base_path = ".".join(path_segments)
+                # When we hit an array, add previous path segments if any
+                if path_segments:
+                    query += "." + ".".join(path_segments)
+                    path_segments = []
+                # Add the array access
                 array_field = part[:-2]
-                remaining_path = ".".join(path_parts[i + 1:])
-
-                if remaining_path:
-                    full_path = (
-                        f"{base_path}.{array_field}" if base_path else array_field
-                    )
-                    return f"$.{full_path}[*] ? ({raw_value})".replace(
-                        "@", f"@.{remaining_path}"
-                    )
-                else:
-                    full_path = (
-                        f"{base_path}.{array_field}" if base_path else array_field
-                    )
-                    return f"$.{full_path}[*] ? ({raw_value})"
+                query += f".{array_field}[*]"
             else:
                 path_segments.append(part)
-    else:
-        # Regular field query
-        return f"$.{field_path} ? ({raw_value})"
+
+        # Add any remaining path segments
+        if path_segments:
+            query += "." + ".".join(path_segments)
+
+        # Add the filter condition
+        query += f" ? ({raw_value})"
+
+        return query
+
+    # Regular field query
+    return f"$.{field_path} ? ({raw_value})"
 
 
 def validate_pipeline_name(pipeline_name: str) -> None:
@@ -185,11 +184,6 @@ def parse_json_filter(filter_str: str) -> tuple:
     pipeline_name, field_spec = parts
     validate_pipeline_name(pipeline_name)
 
-    # Match array queries first
-    # array_match = re.match(r"(.+?)\[\]=(.+)", field_spec)
-    # if array_match:
-    #     return pipeline_name, array_match.group(1), "[]", array_match.group(2)
-
     # Then match regular field queries
     match = re.match(r"(.+?)(~|=|>=|<=|>|<)(.+)", field_spec)
     if not match:

store/neurostore/schemas/data.py

@@ -337,12 +337,27 @@ class BaseStudySchema(BaseDataSchema):
     features = fields.Method("get_features")
 
     def get_features(self, obj):
+        from .pipeline import PipelineStudyResultSchema
+
         pipelines = self.context.get("feature_display", None)
 
         if pipelines is None:
             return {}
 
-        return obj.display_features(pipelines)
+        features = obj.display_features(pipelines)
+        # Flatten each pipeline's predictions
+        if features:
+            flattened_features = {}
+            for pipeline_name, feature_data in features.items():
+                if isinstance(feature_data, dict):
+                    flattened_features[pipeline_name] = (
+                        PipelineStudyResultSchema.flatten_dict(feature_data)
+                    )
+                else:
+                    flattened_features[pipeline_name] = feature_data
+            return flattened_features
+
+        return features
 
     class Meta:
         additional = (

store/neurostore/schemas/pipeline.py

@@ -42,10 +42,43 @@ class PipelineStudyResultSchema(BaseSchema):
     class Meta:
         model = PipelineStudyResult
 
+    @classmethod
+    def flatten_dict(cls, d, parent_key="", sep="."):
+        """Flatten nested dictionaries and arrays containing dictionaries."""
+        items = []
+        for k, v in d.items():
+            new_key = f"{parent_key}{sep}{k}" if parent_key else k
+
+            if isinstance(v, dict):
+                items.extend(cls.flatten_dict(v, new_key, sep=sep).items())
+            elif isinstance(v, list):
+                if v and all(isinstance(item, dict) for item in v):
+                    # For arrays of dictionaries, include index in the key
+                    for idx, dict_item in enumerate(v):
+                        array_key = f"{new_key}[{idx}]"
+                        flattened = cls.flatten_dict(dict_item)
+                        for sub_key, sub_value in flattened.items():
+                            items.append((f"{array_key}.{sub_key}", sub_value))
+                else:
+                    # Keep non-dictionary arrays intact
+                    items.append((new_key, v))
+            else:
+                items.append((new_key, v))
+        return dict(items)
+
     @post_dump
-    def remove_none(self, data, **kwargs):
-        """Remove null values from serialized output."""
-        return {key: value for key, value in data.items() if value is not None}
+    def remove_none_and_flatten(self, data, **kwargs):
+        """Remove null values and flatten nested dictionaries in result_data."""
+        # Remove None values
+        data = {key: value for key, value in data.items() if value is not None}
+
+        # Flatten result_data if it exists
+        if "result_data" in data and isinstance(data["result_data"], dict):
+            # Get predictions section which contains our nested data
+            result_data = data["result_data"]
+            data["result_data"] = self.flatten_dict(result_data)
+
+        return data
 
 
 # Register schemas

store/neurostore/tests/api/test_base_studies.py

@@ -28,11 +28,9 @@ def test_features_query(auth_client, ingest_demographic_features):
     )
     assert result.status_code == 200
     assert "features" in result.json()["results"][0]
-    assert (
-        "age_mean"
-        in result.json()["results"][0]["features"]["ParticipantInfo"]["predictions"][
-            "groups"
-        ][0]
+    features = result.json()["results"][0]["features"]["ParticipantInfo"]
+    assert any(
+        key.startswith("predictions") and key.endswith("].age_mean") for key in features
     )
 
 
@@ -98,8 +96,10 @@ def test_features_query_with_or(auth_client, ingest_demographic_features, sessio
 
     api_diagnoses = set()
     for res in result.json()["results"]:
-        for group in res["features"]["ParticipantInfo"]["predictions"]["groups"]:
-            api_diagnoses.add(group["diagnosis"])
+        features = res["features"]["ParticipantInfo"]
+        # Get all diagnosis values from flattened structure
+        diagnoses = [v for k, v in features.items() if k.endswith(".diagnosis")]
+        api_diagnoses.update(diagnoses)
 
     # Compare database and API results
     assert db_diagnoses == api_diagnoses

store/neurostore/tests/api/test_json_queries.py

@@ -145,6 +145,20 @@ def test_pipeline_multiple_filters(auth_client, study_pipeline_data):
     assert resp.status_code == 200
 
 
+def test_search_list_of_lists(auth_client, study_pipeline_data):
+    """Test search queries on lists of lists."""
+    # Test searching for a specific task name in a list of lists
+    resp = auth_client.get(
+        (
+            "/api/pipeline-study-results?feature_filter="
+            "TaskInfo:predictions.fMRITasks[].Concepts[]~emotion"
+        )
+    )
+    assert resp.status_code == 200
+    results = resp.json()["results"]
+    assert len(results) > 0
+
+
 @pytest.mark.parametrize(
     "query,expected_error",
     [

store/neurostore/tests/api/test_pipeline_resources.py

@@ -90,13 +90,13 @@ def result2(pipeline_study_result_payload, session):
             "TestPipeline:nested.array[]=nested1",
             1,
             "nested1",
-            lambda x: x["result_data"]["nested"]["array"],
+            lambda x: x["result_data"]["nested.array"],
         ),
         (
             "TestPipeline:nested.string~other",
             1,
             "other",
-            lambda x: x["result_data"]["nested"]["string"],
+            lambda x: x["result_data"]["nested.string"],
         ),
         (
             "TestPipeline:array_field[]=value3",

store/neurostore/tests/conftest.py

@@ -702,6 +702,10 @@ def create_pipeline_results(session, ingest_neurosynth, tmp_path):
             "fMRITasks": [
                 {
                     "TaskName": random.choice(["oddball", "n-back", "rest"]),
+                    "Concepts": random.sample(
+                        ["emotion", "memory", "attention", "learning"],
+                        k=random.randint(1, 3),
+                    ),
                     "TaskDescription": (
                         "Participants performed a "
                         f"{random.choice(['visual', 'auditory'])} task"