diff --git a/PWGHF/D2H/Tasks/taskDs.cxx b/PWGHF/D2H/Tasks/taskDs.cxx
index 0614583598f..2621b1e1163 100644
--- a/PWGHF/D2H/Tasks/taskDs.cxx
+++ b/PWGHF/D2H/Tasks/taskDs.cxx
@@ -95,7 +95,7 @@ static std::unordered_map<int8_t, std::unordered_map<int8_t, int8_t>> channelsRe
                                                                                            {Mother::Dplus, {{ResonantChannel::PhiPi, hf_decay::hf_cand_3prong::DecayChannelResonant::DplusToPhiPi}, {ResonantChannel::Kstar0K, hf_decay::hf_cand_3prong::DecayChannelResonant::DplusToKstar0K}}}}};
 
 template <typename T>
-concept hasDsMlInfo = requires(T candidate) {
+concept HasDsMlInfo = requires(T candidate) {
   candidate.mlProbDsToKKPi();
   candidate.mlProbDsToPiKK();
 };
@@ -468,7 +468,7 @@ struct HfTaskDs {
   /// \param candidate is candidate
   /// \param dataType is data class, as defined in DataType enum
   /// \param finalState is either KKPi or PiKK, as defined in FinalState enum
-  template <bool isMc, typename Coll, hasDsMlInfo Cand>
+  template <bool isMc, typename Coll, HasDsMlInfo Cand>
   void fillSparse(const Cand& candidate, DataType dataType, FinalState finalState)
   {
     auto mass = finalState == FinalState::KKPi ? hfHelper.invMassDsToKKPi(candidate) : hfHelper.invMassDsToPiKK(candidate);
@@ -596,13 +596,10 @@ struct HfTaskDs {
   {
 
     int id = o2::constants::physics::Pdg::kDS;
-    auto yCand = hfHelper.yDs(candidate);
     if (dataType == DataType::McDplusPrompt || dataType == DataType::McDplusNonPrompt || dataType == DataType::McDplusBkg) {
       id = o2::constants::physics::Pdg::kDPlus;
-      yCand = hfHelper.yDplus(candidate);
     } else if (dataType == DataType::McLcBkg) {
       id = o2::constants::physics::Pdg::kLambdaCPlus;
-      yCand = hfHelper.yLc(candidate);
     }
 
     auto indexMother = RecoDecay::getMother(mcParticles,
@@ -610,13 +607,14 @@ struct HfTaskDs {
                                             id, true);
 
     if (indexMother != -1) {
-      if (yCandRecoMax >= 0. && std::abs(yCand) > yCandRecoMax) {
-        return;
-      }
 
       auto pt = candidate.pt(); // rec. level pT
 
       if (candidate.isSelDsToKKPi() >= selectionFlagDs) { // KKPi
+        auto yCand = candidate.y(hfHelper.invMassDsToKKPi(candidate));
+        if (yCandRecoMax >= 0. && std::abs(yCand) > yCandRecoMax) {
+          return;
+        }
         fillHisto(candidate, dataType);
         fillHistoKKPi<true, Coll>(candidate, dataType);
 
@@ -631,6 +629,10 @@ struct HfTaskDs {
         }
       }
       if (candidate.isSelDsToPiKK() >= selectionFlagDs) { // PiKK
+        auto yCand = candidate.y(hfHelper.invMassDsToPiKK(candidate));
+        if (yCandRecoMax >= 0. && std::abs(yCand) > yCandRecoMax) {
+          return;
+        }
         fillHisto(candidate, dataType);
         fillHistoPiKK<true, Coll>(candidate, dataType);
 
@@ -651,15 +653,17 @@ struct HfTaskDs {
   template <typename Coll, typename CandDs>
   void runDataAnalysisPerCandidate(CandDs const& candidate)
   {
-    if (yCandRecoMax >= 0. && std::abs(hfHelper.yDs(candidate)) > yCandRecoMax) {
-      return;
-    }
-
     if (candidate.isSelDsToKKPi() >= selectionFlagDs) { // KKPi
+      if (yCandRecoMax >= 0. && std::abs(candidate.y(hfHelper.invMassDsToKKPi(candidate))) > yCandRecoMax) {
+        return;
+      }
       fillHisto(candidate, DataType::Data);
       fillHistoKKPi<false, Coll>(candidate, DataType::Data);
     }
     if (candidate.isSelDsToPiKK() >= selectionFlagDs) { // PiKK
+      if (yCandRecoMax >= 0. && std::abs(candidate.y(hfHelper.invMassDsToPiKK(candidate))) > yCandRecoMax) {
+        return;
+      }
       fillHisto(candidate, DataType::Data);
       fillHistoPiKK<false, Coll>(candidate, DataType::Data);
     }
@@ -687,16 +691,19 @@ struct HfTaskDs {
       }
     }
     if (isBkg && fillMcBkgHistos) {
-      if (yCandRecoMax >= 0. && std::abs(hfHelper.yDs(candidate)) > yCandRecoMax) {
-        return;
-      }
 
       if (candidate.isSelDsToKKPi() >= selectionFlagDs || candidate.isSelDsToPiKK() >= selectionFlagDs) {
         if (candidate.isSelDsToKKPi() >= selectionFlagDs) { // KKPi
+          if (yCandRecoMax >= 0. && std::abs(candidate.y(hfHelper.invMassDsToKKPi(candidate))) > yCandRecoMax) {
+            return;
+          }
           fillHisto(candidate, DataType::McBkg);
           fillHistoKKPi<true, Coll>(candidate, DataType::McBkg);
         }
         if (candidate.isSelDsToPiKK() >= selectionFlagDs) { // PiKK
+          if (yCandRecoMax >= 0. && std::abs(candidate.y(hfHelper.invMassDsToPiKK(candidate))) > yCandRecoMax) {
+            return;
+          }
           fillHisto(candidate, DataType::McBkg);
           fillHistoPiKK<true, Coll>(candidate, DataType::McBkg);
         }
diff --git a/PWGHF/TableProducer/treeCreatorDsToKKPi.cxx b/PWGHF/TableProducer/treeCreatorDsToKKPi.cxx
index bfbd0b38ea3..6cff1ecbb06 100644
--- a/PWGHF/TableProducer/treeCreatorDsToKKPi.cxx
+++ b/PWGHF/TableProducer/treeCreatorDsToKKPi.cxx
@@ -324,42 +324,40 @@ struct HfTreeCreatorDsToKKPi {
   template <bool doMc = false, int massHypo = 0, typename Coll, typename T>
   void fillCandidateTable(const T& candidate)
   {
-    int8_t flagMc = 0;
-    int8_t originMc = 0;
-    int8_t channelMc = 0;
-    int8_t isSwapped = massHypo; // 0 if KKPi, 1 if PiKK
-    float yCand = 0;
-    float eCand = 0;
-    float ctCand = 0;
+    float invMassDs = 0;
+    float deltaMassPhiKK = 0;
+    float absCos3PiKDs = 0;
+    if constexpr (massHypo == 0) {
+      invMassDs = hfHelper.invMassDsToKKPi(candidate);
+      deltaMassPhiKK = hfHelper.deltaMassPhiDsToKKPi(candidate);
+      absCos3PiKDs = hfHelper.absCos3PiKDsToKKPi(candidate);
+    } else if constexpr (massHypo == 1) {
+      invMassDs = hfHelper.invMassDsToPiKK(candidate);
+      deltaMassPhiKK = hfHelper.deltaMassPhiDsToPiKK(candidate);
+      absCos3PiKDs = hfHelper.absCos3PiKDsToPiKK(candidate);
+    }
+
+    int8_t flagMc{0};
+    int8_t originMc{0};
+    int8_t channelMc{0};
+    int8_t isSwapped{massHypo}; // 0 if KKPi, 1 if PiKK
+    float eCand{0.f};
+    float ctCand{0.f};
+    float yCand = candidate.y(invMassDs);
     if constexpr (doMc) {
       flagMc = candidate.flagMcMatchRec();
       originMc = candidate.originMcRec();
       channelMc = candidate.flagMcDecayChanRec();
       isSwapped = candidate.isCandidateSwapped();
       if (fillDplusMc && candidate.flagMcDecayChanRec() == channelsResonant[Mother::Dplus][decayChannel]) {
-        yCand = hfHelper.yDplus(candidate);
         eCand = hfHelper.eDplus(candidate);
         ctCand = hfHelper.ctDplus(candidate);
       } else {
-        yCand = hfHelper.yDs(candidate);
         eCand = hfHelper.eDs(candidate);
         ctCand = hfHelper.ctDs(candidate);
       }
     }
 
-    float invMassDs = 0;
-    float deltaMassPhiKK = 0;
-    float absCos3PiKDs = 0;
-    if constexpr (massHypo == 0) {
-      invMassDs = hfHelper.invMassDsToKKPi(candidate);
-      deltaMassPhiKK = hfHelper.deltaMassPhiDsToKKPi(candidate);
-      absCos3PiKDs = hfHelper.absCos3PiKDsToKKPi(candidate);
-    } else if constexpr (massHypo == 1) {
-      invMassDs = hfHelper.invMassDsToPiKK(candidate);
-      deltaMassPhiKK = hfHelper.deltaMassPhiDsToPiKK(candidate);
-      absCos3PiKDs = hfHelper.absCos3PiKDsToPiKK(candidate);
-    }
-
     auto const& collision = candidate.template collision_as<Coll>();
     float centrality = o2::hf_centrality::getCentralityColl(collision);
 
@@ -617,7 +615,7 @@ struct HfTreeCreatorDsToKKPi {
         particle.pt(),
         particle.eta(),
         particle.phi(),
-        RecoDecay::y(particle.pVector(), o2::constants::physics::MassDS),
+        std::abs(particle.pdgCode()) == o2::constants::physics::Pdg::kDS ? RecoDecay::y(particle.pVector(), o2::constants::physics::MassDS) : RecoDecay::y(particle.pVector(), o2::constants::physics::MassDPlus),
         particle.flagMcMatchGen(),
         particle.originMcGen());
     }