Merge pull request #949 from astronomerritt/indexing_sql

Adding indexes to SQL outputs.

src/sorcha/modules/PPOutput.py

@@ -67,7 +67,7 @@ def PPOutWriteHDF5(pp_results, outf, keyin):
     return of
 
 
-def PPOutWriteSqlite3(pp_results, outf):
+def PPOutWriteSqlite3(pp_results, outf, lastchunk=False, tablename="sorcha_results"):
     """
     Writes a pandas dataframe out to a CSV file at a location given by the user.
 
@@ -90,12 +90,22 @@ def PPOutWriteSqlite3(pp_results, outf):
 
     cnx = sqlite3.connect(outf)
 
-    pp_results.to_sql("sorcha_results", con=cnx, if_exists="append", index=False)
+    pp_results.to_sql(tablename, con=cnx, if_exists="append", index=False)
 
-    pplogger.info("SQL results saved in table sorcha_results in database {}.".format(outf))
+    # we don't want to index the table until we're sure we're done appending to it
+    # as recreating the indexes on every append is slow
+    if lastchunk:
+        pplogger.info("Last chunk detected. Indexing SQL table...")
+        cur = cnx.cursor()
+        cur.execute("CREATE INDEX ObjID ON {} (ObjID)".format(tablename))
+        cur.execute("CREATE INDEX fieldMJD_TAI ON {} (fieldMJD_TAI)".format(tablename))
+        cur.execute("CREATE INDEX optFilter ON {} (optFilter)".format(tablename))
+        cnx.commit()
 
+    pplogger.info("SQL results saved in table {} in database {}.".format(tablename, outf))
 
-def PPWriteOutput(cmd_args, configs, observations_in, endChunk=0, verbose=False):
+
+def PPWriteOutput(cmd_args, configs, observations_in, endChunk=0, verbose=False, lastchunk=False):
     """
     Writes the output in the format specified in the config file to a location
     specified by the user.
@@ -214,7 +224,7 @@ def PPWriteOutput(cmd_args, configs, observations_in, endChunk=0, verbose=False)
         outputsuffix = ".db"
         out = os.path.join(cmd_args.outpath, cmd_args.outfilestem + outputsuffix)
         verboselog("Output to sqlite3 database...")
-        observations = PPOutWriteSqlite3(observations, out)
+        observations = PPOutWriteSqlite3(observations, out, lastchunk)
 
     elif configs["output_format"] == "hdf5" or configs["output_format"] == "h5":
         outputsuffix = ".h5"

src/sorcha/sorcha.py

@@ -161,6 +161,7 @@ def runLSSTSimulation(args, configs):
     startChunk = 0
     endChunk = 0
     loopCounter = 0
+    lastChunk = False
 
     ii = -1
     with open(args.orbinfile) as f:
@@ -178,6 +179,9 @@ def runLSSTSimulation(args, configs):
         endChunk = startChunk + configs["size_serial_chunk"]
         verboselog("Working on objects {}-{}".format(startChunk, endChunk))
 
+        if endChunk >= lenf:
+            lastChunk = True
+
         # Processing begins, all processing is done for chunks
         if configs["ephemerides_type"].casefold() == "external":
             verboselog("Reading in chunk of orbits and associated ephemeris from an external file")
@@ -334,7 +338,7 @@ def runLSSTSimulation(args, configs):
         pplogger.info("Output results for this chunk")
 
         # write output
-        PPWriteOutput(args, configs, observations, endChunk, verbose=args.verbose)
+        PPWriteOutput(args, configs, observations, endChunk, verbose=args.verbose, lastchunk=lastChunk)
 
         startChunk = startChunk + configs["size_serial_chunk"]
         loopCounter = loopCounter + 1

tests/sorcha/test_PPOutput.py

@@ -121,7 +121,7 @@ def test_PPWriteOutput_sql(tmp_path):
         dtype=object,
     )
 
-    PPWriteOutput(args, configs, observations, 10)
+    PPWriteOutput(args, configs, observations, 10, lastchunk=True)
     cnx = sqlite3.connect(os.path.join(tmp_path, "PPOutput_test_out.db"))
     cur = cnx.cursor()
     cur.execute("select * from sorcha_results")
@@ -130,6 +130,13 @@ def test_PPWriteOutput_sql(tmp_path):
 
     assert_equal(sql_test_in.loc[0, :].values, expected)
 
+    # check indexes were properly created
+    cur.execute("PRAGMA index_list('sorcha_results')")
+    indexes = cur.fetchall()
+
+    index_list = [indexes[i][1] for i in range(0, 3)]
+    assert index_list == ["optFilter", "fieldMJD_TAI", "ObjID"]
+
 
 def test_PPWriteOutput_all(tmp_path):
     # additional test to ensure that "all" output option and no rounding works