Merge pull request #298 from punch-mission/querying-work

Speed up some DB queries

Author: jmbhughes

punchpipe/control/db.py

@@ -61,11 +61,11 @@ def directory(self, root: str):
         return os.path.join(root, self.level, self.file_type + self.observatory, self.date_obs.strftime("%Y/%m/%d"))
 
 
-Index("date_obs_index", File.date_obs, mysql_using="btree", mariadb_using="btree")
-Index("observatory_index", File.observatory, mysql_using="hash", mariadb_using="hash")
-Index("file_type_index", File.file_type, mysql_using="hash", mariadb_using="hash")
-Index("processing_flow_index", File.processing_flow, mysql_using="hash", mariadb_using="hash")
-Index("ready_files_index", File.file_type, File.level, File.state, File.observatory, File.outlier, File.date_obs)
+Index("get_ready_files", File.state, File.level, File.date_obs.desc(), File.file_type, File.observatory)
+Index("get_ready_files_alt", File.level, File.date_obs.desc(), File.file_type, File.state)
+Index("construct_background", File.level, File.observatory, File.outlier, File.date_obs, File.state, File.file_type)
+Index("get_cal_file", File.file_type, File.observatory, File.date_obs, File.state)
+Index("CNN", File.file_type, File.observatory, File.level, File.state, File.outlier)
 
 
 class Flow(Base):
@@ -88,6 +88,10 @@ def __repr__(self):
         return f"Flow(id={self.flow_id!r})"
 
 
+Index("flow_stats", Flow.end_time, Flow.flow_type, Flow.state)
+Index("flow_cards", Flow.start_time, Flow.flow_level, Flow.flow_type)
+
+
 class FileRelationship(Base):
     __tablename__ = "relationships"
     relationship_id = Column(Integer, primary_key=True)

punchpipe/flows/level1.py

@@ -223,15 +223,10 @@ def get_two_closest_stray_light(level0_file, session=None, max_distance: timedel
                         text("second"), File.date_obs, level0_file.date_obs)))
                   .filter(File.file_type == model_type)
                   .filter(File.observatory == level0_file.observatory)
-                  .filter(File.state == "created")
-                  .filter(File.file_version.not_like("v%"))) #filters out "v0a".
+                  .filter(File.state == "created"))
     if max_distance:
         best_models = best_models.filter(dt < max_distance.total_seconds())
-    highest_version = best_models.order_by(File.file_version).first()
-    if highest_version is None:
-        return None, None
-    highest_version = highest_version[0].file_version
-    best_models = best_models.filter(File.file_version == highest_version).order_by(dt.asc()).limit(2).all()
+    best_models = best_models.order_by(dt.asc()).limit(2).all()
     if len(best_models) < 2:
         return None, None
     # Drop the dt values
@@ -246,17 +241,13 @@ def get_two_best_stray_light(level0_file, session=None):
     before_model = (session.query(File)
                     .filter(File.file_type == model_type)
                     .filter(File.observatory == level0_file.observatory)
-                    .filter(File.level == '1')
                     .filter(File.date_obs < level0_file.date_obs)
-                    .filter(File.file_version.not_like("v%")) #filters out "v0a".
-                    .order_by(File.file_version.desc(), File.date_obs.desc()).first())
+                    .order_by(File.date_obs.desc()).first())
     after_model = (session.query(File)
                    .filter(File.file_type == model_type)
                    .filter(File.observatory == level0_file.observatory)
-                   .filter(File.level == '1')
                    .filter(File.date_obs > level0_file.date_obs)
-                   .filter(File.file_version.not_like("v%")) #filters out "v0a".
-                   .order_by(File.file_version.desc(), File.date_obs.asc()).first())
+                   .order_by(File.date_obs.asc()).first())
     if before_model is None or after_model is None:
         # We're waiting for the scheduler to fill in here and tell us what's what
         return None, None
@@ -285,6 +276,15 @@ def get_two_best_stray_light(level0_file, session=None):
     return None, None
 
 
+def get_first_last_stray_light(session):
+    dates = (session.query(func.min(File.date_obs), func.max(File.date_obs))
+             .where(File.file_type.like('S%')).
+             where(File.state == 'created')).all()
+    if len(dates) == 0:
+        return datetime(1900, 1, 1), datetime(2900, 1, 1)
+    return dates[0]
+
+
 def get_quartic_model_paths(level0_files, pipeline_config: dict, session=None):
     # Get all models, in reverse-chronological order
     models = (session.query(File)
@@ -446,6 +446,7 @@ def level1_early_process_flow(flow_id: int | list[int], pipeline_config_path=Non
 @task(cache_policy=NO_CACHE)
 def level1_late_query_ready_files(session, pipeline_config: dict, reference_time=None, max_n=9e99):
     logger = get_run_logger()
+    start_date, end_date = get_first_last_stray_light(session)
     parent = aliased(File)
     child = aliased(File)
     child_exists_subquery = (session.query(parent)
@@ -459,6 +460,8 @@ def level1_late_query_ready_files(session, pipeline_config: dict, reference_time
              .filter(File.level == "1")
              .filter(File.state.in_(["created", "progressed"]))
              .filter(~child_exists_subquery)
+             .filter(File.date_obs >= start_date)
+             .filter(File.date_obs <= end_date)
              .order_by(File.date_obs.desc()).all())
 
     distortion_paths = get_distortion_paths(ready, pipeline_config, session)

punchpipe/flows/tests/test_level1.py

@@ -128,47 +128,31 @@ def session_fn(session):
                        file_type="SM",
                        observatory='2',
                        state='created',
-                       file_version='0b',
+                       file_version='1',
                        software_version='none',
                        date_obs=datetime.now(UTC)-timedelta(hours=12))
 
     stray_light_before1 = File(level="1",
                        file_type="SM",
                        observatory='2',
                        state='created',
-                       file_version='1',
+                       file_version='2',
                        software_version='none',
                        date_obs=datetime.now(UTC)-timedelta(hours=16))
 
-    stray_light_before0 = File(level="1",
-                       file_type="SM",
-                       observatory='2',
-                       state='created',
-                       file_version='0b',
-                       software_version='none',
-                       date_obs=datetime.now(UTC)-timedelta(hours=12))
-
-    stray_light_after2 = File(level="1",
-                       file_type="SM",
-                       observatory='2',
-                       state='created',
-                       file_version='0b',
-                       software_version='none',
-                       date_obs=datetime.now(UTC)+timedelta(hours=12))
-
     stray_light_after1 = File(level="1",
                        file_type="SM",
                        observatory='2',
                        state='created',
-                       file_version='1',
+                       file_version='2',
                        software_version='none',
                        date_obs=datetime.now(UTC)+timedelta(hours=16))
 
     stray_light_after0 = File(level="1",
                        file_type="SM",
                        observatory='2',
                        state='created',
-                       file_version='0b',
+                       file_version='1',
                        software_version='none',
                        date_obs=datetime.now(UTC)+timedelta(hours=12))
 
@@ -236,12 +220,10 @@ def session_fn(session):
     session.add(distortion0)
     session.add(distortion1)
     session.add(distortion2)
-    session.add(stray_light_before0)
     session.add(stray_light_before1)
     session.add(stray_light_before2)
     session.add(stray_light_after0)
     session.add(stray_light_after1)
-    session.add(stray_light_after2)
     session.add(mask_file0)
     session.add(mask_file1)
     session.add(mask_file2)