Tweak carton/program query with specific settings

python/valis/db/queries.py

@@ -4,6 +4,7 @@
 
 # all resuable queries go here
 
+from contextlib import contextmanager
 import itertools
 import packaging
 import uuid
@@ -61,16 +62,15 @@ def append_pipes(query: peewee.ModelSelect, table: str = 'stacked',
     # Run initial query as a temporary table.
     temp = create_temporary_table(query, indices=['sdss_id'])
 
-    model = vizdb.SDSSidStacked if table == 'stacked' else vizdb.SDSSidFlat
     qq = temp.select(temp.__star__,
-                     vizdb.SDSSidToPipes.in_boss,
-                     vizdb.SDSSidToPipes.in_apogee,
-                     vizdb.SDSSidToPipes.in_bvs,
-                     vizdb.SDSSidToPipes.in_astra,
-                     vizdb.SDSSidToPipes.has_been_observed,
-                     vizdb.SDSSidToPipes.release,
-                     vizdb.SDSSidToPipes.obs,
-                     vizdb.SDSSidToPipes.mjd).\
+                    vizdb.SDSSidToPipes.in_boss,
+                    vizdb.SDSSidToPipes.in_apogee,
+                    vizdb.SDSSidToPipes.in_bvs,
+                    vizdb.SDSSidToPipes.in_astra,
+                    vizdb.SDSSidToPipes.has_been_observed,
+                    vizdb.SDSSidToPipes.release,
+                    vizdb.SDSSidToPipes.obs,
+                    vizdb.SDSSidToPipes.mjd).\
         join(vizdb.SDSSidToPipes, on=(temp.c.sdss_id == vizdb.SDSSidToPipes.sdss_id)).\
         distinct(temp.c.sdss_id)
 
@@ -292,7 +292,14 @@ def carton_program_search(name: str,
     """
 
     if query is None:
-        query = vizdb.SDSSidStacked.select(vizdb.SDSSidStacked)
+        query = vizdb.SDSSidStacked.select(vizdb.SDSSidStacked).distinct()
+
+    # NOTE: These setting seem to help when querying some cartons or programs, mainly
+    # those with small number of targets, and in some cases with these the query
+    # actually applies the LIMIT more efficiently, but it's not a perfect solution.
+    vizdb.database.execute_sql('SET enable_gathermerge = off;')
+    vizdb.database.execute_sql('SET parallel_tuple_cost = 100;')
+    vizdb.database.execute_sql('SET enable_bitmapscan = off;')
 
     query = (query.join(
                 vizdb.SDSSidFlat,
@@ -944,7 +951,8 @@ def get_target_by_altid(id: str | int, idtype: str = None) -> peewee.ModelSelect
     return get_targets_by_sdss_id(res.sdss_id)
 
 
-def create_temporary_table(query: peewee.ModelSelect, indices: list[str] | None = None) -> peewee.Table:
+def create_temporary_table(query: peewee.ModelSelect,
+                           indices: list[str] | None = None) -> Generator[None, None, peewee.Table]:
     """Create a temporary table from a query."""
 
     table_name = uuid.uuid4().hex[0:8]

python/valis/routes/query.py

@@ -41,6 +41,7 @@ class SearchModel(BaseModel):
     program: Optional[str] = Field(None, description='The program name', example='bhm_rm')
     carton: Optional[str] = Field(None, description='The carton name', example='bhm_rm_core')
     observed: Optional[bool] = Field(True, description='Flag to only include targets that have been observed', example=True)
+    limit: Optional[int] = Field(None, description='Limit the number of returned targets', example=100)
 
 class MainResponse(SDSSModel):
     """ Combined model from all individual query models """
@@ -105,6 +106,13 @@ async def main_search(self, body: SearchModel):
             query = carton_program_search(body.program or body.carton,
                                           'program' if body.program else 'carton',
                                           query=query)
+
+        # DANGER!!! This limit applies *before* the append_pipes call. If the
+        # append_pipes call includes observed=True we may have limited things in
+        # such a way that only unobserved or very few targets are returned.
+        if body.limit:
+            query = query.limit(body.limit)
+
         # append query to pipes
         if query:
             query = append_pipes(query, observed=body.observed)