update reserved size

Author: zjxiongOvO

PWGDQ/Tasks/dqEfficiency_withAssoc.cxx

@@ -2005,23 +2005,34 @@ struct AnalysisSameEventPairing {
     uint32_t mcDecision = static_cast<uint32_t>(0);
     bool isCorrectAssoc_leg1 = false;
     bool isCorrectAssoc_leg2 = false;
-    dielectronList.reserve(1);
-    dimuonList.reserve(1);
-    dielectronsExtraList.reserve(1);
-    dimuonsExtraList.reserve(1);
-    dielectronInfoList.reserve(1);
-    dileptonInfoList.reserve(1);
+
+    // estimate reserved size
+    int64_t reserveSize = 0;
+    for (auto& event : events) {
+      if (event.isEventSelected_bit(0)) {
+        auto groupedAssocs = assocs.sliceBy(preslice, event.globalIndex());
+        reserveSize += (groupedAssocs.size() * (groupedAssocs.size() - 1)) / 2; // n choose 2 combinations
+      }
+    }
+
+    dielectronList.reserve(reserveSize);
+    dimuonList.reserve(reserveSize);
+    dielectronsExtraList.reserve(reserveSize);
+    dimuonsExtraList.reserve(reserveSize);
+    dielectronInfoList.reserve(reserveSize);
+    dileptonInfoList.reserve(reserveSize);
     if (fConfigOptions.flatTables.value) {
-      dielectronAllList.reserve(1);
-      dimuonAllList.reserve(1);
+      dielectronAllList.reserve(reserveSize);
+      dimuonAllList.reserve(reserveSize);
     }
     if (useMiniTree.fConfigMiniTree) {
-      dileptonMiniTreeGen.reserve(1);
-      dileptonMiniTreeRec.reserve(1);
+      dileptonMiniTreeGen.reserve(reserveSize);
+      dileptonMiniTreeRec.reserve(reserveSize);
     }
     if (fConfigOptions.polarTables.value) {
-      dileptonPolarList.reserve(1);
+      dileptonPolarList.reserve(reserveSize);
     }
+
     constexpr bool eventHasQvector = ((TEventFillMap & VarManager::ObjTypes::ReducedEventQvector) > 0);
     constexpr bool trackHasCov = ((TTrackFillMap & VarManager::ObjTypes::ReducedTrackBarrelCov) > 0);
 

PWGDQ/Tasks/dqEfficiency_withAssoc_direct.cxx

@@ -1691,20 +1691,6 @@ struct AnalysisSameEventPairing {
       dileptonMiniTreeRec.reserve(reserveSize);
     }
 
-    // dielectronList.reserve(1);
-    // // dimuonList.reserve(1);
-    // dielectronsExtraList.reserve(1);
-    // // dimuonsExtraList.reserve(1);
-    // dielectronInfoList.reserve(1);
-    // dileptonInfoList.reserve(1);
-    // if (fConfigOptions.flatTables.value) {
-    //   dielectronAllList.reserve(1);
-    //   // dimuonAllList.reserve(1);
-    // }
-    // if (fConfigOptions.fConfigMiniTree) {
-    //   dileptonMiniTreeGen.reserve(1);
-    //   dileptonMiniTreeRec.reserve(1);
-    // }
     constexpr bool eventHasQvector = ((TEventFillMap & VarManager::ObjTypes::CollisionQvect) > 0);
     constexpr bool trackHasCov = ((TTrackFillMap & VarManager::ObjTypes::TrackCov) > 0);
 

PWGDQ/Tasks/tableReader_withAssoc.cxx

@@ -1829,22 +1829,32 @@ struct AnalysisSameEventPairing {
     // Reserve capacity for the output tables to avoid repeated reallocations
     // inside the Arrow builders.  Unused capacity is virtual address space
     // only — pages are not faulted in until written.
-    auto nAssocs = assocs.size();
-    dielectronList.reserve(nAssocs);
-    dimuonList.reserve(nAssocs);
-    dielectronsExtraList.reserve(nAssocs);
-    dielectronInfoList.reserve(nAssocs);
-    dimuonsExtraList.reserve(nAssocs);
-    dileptonInfoList.reserve(nAssocs);
-    dileptonFlowList.reserve(nAssocs);
+    // estimate reserved size
+    int64_t reserveSize = 0;
+    for (auto& event : events) {
+      if (event.isEventSelected_bit(0)) {
+        auto groupedAssocs = assocs.sliceBy(preslice, event.globalIndex());
+        reserveSize += (groupedAssocs.size() * (groupedAssocs.size() - 1)) / 2; // n choose 2 combinations
+      }
+    }
+    LOG(info) << "Reserving capacity for " << reserveSize << " pairs in the output tables";
+    
+    dielectronList.reserve(reserveSize);
+    dimuonList.reserve(reserveSize);
+    dielectronsExtraList.reserve(reserveSize);
+    dielectronInfoList.reserve(reserveSize);
+    dimuonsExtraList.reserve(reserveSize);
+    dileptonInfoList.reserve(reserveSize);
+    dileptonFlowList.reserve(reserveSize);
     if (fConfigOptions.flatTables.value) {
-      dielectronAllList.reserve(nAssocs);
-      dimuonAllList.reserve(nAssocs);
+      dielectronAllList.reserve(reserveSize);
+      dimuonAllList.reserve(reserveSize);
     }
     if (fConfigOptions.polarTables.value) {
-      dileptonPolarList.reserve(nAssocs);
-      dileptonEventInfoList.reserve(nAssocs);
+      dileptonPolarList.reserve(reserveSize);
+      dileptonEventInfoList.reserve(reserveSize);
     }
+    
     fAmbiguousPairs.clear();
     constexpr bool eventHasQvector = ((TEventFillMap & VarManager::ObjTypes::ReducedEventQvector) > 0);
     constexpr bool eventHasQvectorCentr = ((TEventFillMap & VarManager::ObjTypes::CollisionQvect) > 0);

PWGDQ/Tasks/tableReader_withAssoc_direct.cxx

@@ -1637,16 +1637,21 @@ struct AnalysisSameEventPairing {
     int sign1 = 0;
     int sign2 = 0;
 
-    dielectronList.reserve(1);
-    // dimuonList.reserve(1);
-    dielectronsExtraList.reserve(1);
-    // dimuonsExtraList.reserve(1);
-    dielectronInfoList.reserve(1);
-    dileptonInfoList.reserve(1);
+    // estimate reserved size
+    int64_t reserveSize = 0;
+    for (auto& event : events) {
+      if (event.isEventSelected_bit(0)) {
+        auto groupedAssocs = assocs.sliceBy(preslice, event.globalIndex());
+        reserveSize += (groupedAssocs.size() * (groupedAssocs.size() - 1)) / 2; // n choose 2 combinations
+      }
+    }
 
+    dielectronList.reserve(reserveSize);
+    dielectronsExtraList.reserve(reserveSize);
+    dielectronInfoList.reserve(reserveSize);
+    dileptonInfoList.reserve(reserveSize);
     if (fConfigOptions.flatTables.value) {
-      dielectronAllList.reserve(1);
-      // dimuonAllList.reserve(1);
+      dielectronAllList.reserve(reserveSize);
     }
 
     constexpr bool eventHasQvector = ((TEventFillMap & VarManager::ObjTypes::CollisionQvect) > 0);