Write to the APDB in chunks to avoid overwhelming the database

Author: isullivan

python/lsst/ap/association/diaPipe.py

@@ -309,6 +309,13 @@ class DiaPipelineConfig(pipeBase.PipelineTaskConfig,
         doc="Pad the image by this many pixels before removing off-image "
             "diaObjects for association.",
     )
+    maximumTableLength = pexConfig.Field(
+        dtype=int,
+        default=65535,
+        doc="Maximum length of tables allowed to be written in one operation"
+            " to the Cassandra APDB. The default is the hard maximum, "
+            "but performance may be improved with lower values",
+    )
     idGenerator = DetectorVisitIdGeneratorConfig.make_field()
 
     def setDefaults(self):
@@ -782,12 +789,35 @@ def writeToApdb(self, updatedDiaObjects, associatedDiaSources, diaForcedSources)
         diaObjectStore = dropEmptyColumns(self.schema, updatedDiaObjects, tableName="DiaObject")
         diaSourceStore = dropEmptyColumns(self.schema, associatedDiaSources, tableName="DiaSource")
         diaForcedSourceStore = dropEmptyColumns(self.schema, diaForcedSources, tableName="DiaForcedSource")
-        self.apdb.store(
-            DateTime.now().toAstropy(),
-            diaObjectStore,
-            diaSourceStore,
-            diaForcedSourceStore)
-        self.log.info("APDB updated.")
+
+        nObject = len(diaObjectStore)
+        nSource = len(diaSourceStore)
+        nForced = len(diaForcedSourceStore)
+        nObjWritten = 0
+        start = 0
+        end = min(nObject - 1, self.config.maximumTableLength - 1)
+        time = DateTime.now().toAstropy()
+        while nObjWritten < nObject:
+            srcEnd = min(start + self.config.maximumTableLength, nSource)
+            if srcEnd <= start:
+                finalDiaSources = None
+            else:
+                finalDiaSources = diaSourceStore.iloc[start:srcEnd]
+
+            fSrcEnd = min(start + self.config.maximumTableLength, nForced)
+            if fSrcEnd <= start:
+                finalForcedSources = None
+            else:
+                finalForcedSources = diaForcedSourceStore.iloc[start:fSrcEnd]
+            self.apdb.store(
+                time,
+                diaObjectStore.iloc[start:end],
+                finalDiaSources,
+                finalForcedSources)
+            self.log.info("APDB updated.")
+            start = end + 1
+            nObjWritten += self.config.maximumTableLength
+            end = np.min(nObject - 1, start + self.config.maximumTableLength - 1)
 
     def testDataFrameIndex(self, df):
         """Test the sorted DataFrame index for duplicates.