Skip to content

Commit

Permalink
TBranch::GetBulkEntries proper handling of already loaded basket
Browse files Browse the repository at this point in the history
This fixes Issue root-project#6416 and Issue root-project#6417
  • Loading branch information
pcanal authored and ktf committed Sep 30, 2020
1 parent 90fc8ca commit fe9a8a6
Showing 1 changed file with 66 additions and 23 deletions.
89 changes: 66 additions & 23 deletions tree/tree/src/TBranch.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -1322,6 +1322,10 @@ const char* TBranch::GetIconName() const
/// the user_buffer will back the memory of the newly-constructed basket.
///
/// Assumes that this branch is enabled.
///
/// Returns -1 if the entry does not exist
/// Returns -2 in case of error
/// Returns the index of the basket in case of success.
Int_t TBranch::GetBasketAndFirst(TBasket *&basket, Long64_t &first,
TBuffer *user_buffer)
{
Expand All @@ -1332,9 +1336,10 @@ Int_t TBranch::GetBasketAndFirst(TBasket *&basket, Long64_t &first,
// make sure basket buffers are in memory.
basket = fCurrentBasket;
first = fFirstBasketEntry;
return fReadBasket;
} else {
if ((entry < fFirstEntry) || (entry >= fEntryNumber)) {
return 0;
return -1;
}
first = fFirstBasketEntry;
Long64_t last = fNextBasketEntry - 1;
Expand All @@ -1344,7 +1349,7 @@ Int_t TBranch::GetBasketAndFirst(TBasket *&basket, Long64_t &first,
if (fReadBasket < 0) {
fNextBasketEntry = -1;
Error("GetBasketAndFirst", "In the branch %s, no basket contains the entry %lld\n", GetName(), entry);
return -1;
return -2;
}
if (fReadBasket == fWriteBasket) {
fNextBasketEntry = fEntryNumber;
Expand All @@ -1363,7 +1368,7 @@ Int_t TBranch::GetBasketAndFirst(TBasket *&basket, Long64_t &first,
fCurrentBasket = 0;
fFirstBasketEntry = -1;
fNextBasketEntry = -1;
return -1;
return -2;
}
if (fTree->GetClusterPrefetch()) {
TTree::TClusterIterator clusterIterator = fTree->GetClusterIterator(entry);
Expand All @@ -1377,18 +1382,18 @@ Int_t TBranch::GetBasketAndFirst(TBasket *&basket, Long64_t &first,
// cause a reset of the first / next basket entries back to -1.
fFirstBasketEntry = first;
fNextBasketEntry = updatedNext;
}
if (user_buffer) {
// Disassociate basket from memory buffer for bulk IO
// When the user provides a memory buffer (i.e., for bulk IO), we should
// make sure to drop all references to that buffer in the TTree afterward.
fCurrentBasket = nullptr;
fBaskets[fReadBasket] = nullptr;
if (user_buffer) {
// Disassociate basket from memory buffer for bulk IO
// When the user provides a memory buffer (i.e., for bulk IO), we should
// make sure to drop all references to that buffer in the TTree afterward.
fCurrentBasket = nullptr;
fBaskets[fReadBasket] = nullptr;
}
} else {
fCurrentBasket = basket;
}
return fReadBasket;
}
return 1;
}

////////////////////////////////////////////////////////////////////////////////
Expand Down Expand Up @@ -1437,7 +1442,7 @@ Int_t TBranch::GetBulkEntries(Long64_t entry, TBuffer &user_buf)
TBasket *basket = nullptr;
Long64_t first;
Int_t result = GetBasketAndFirst(basket, first, &user_buf);
if (R__unlikely(result <= 0)) return -1;
if (R__unlikely(result < 0)) return -1;
// Only support reading from full clusters.
if (R__unlikely(entry != first)) {
//printf("Failed to read from full cluster; first entry is %ld; requested entry is %ld.\n", first, entry);
Expand All @@ -1458,22 +1463,41 @@ Int_t TBranch::GetBulkEntries(Long64_t entry, TBuffer &user_buf)
return -1;
}

if (&user_buf != buf) {
// The basket was already in memory and might (and might not) be backed by persistent
// storage.
R__ASSERT(result == fReadBasket);
if (fBasketSeek[fReadBasket]) {
// It is backed, so we can be destructive
user_buf.SetBuffer(buf->Buffer(), buf->BufferSize());
buf->ResetBit(TBufferIO::kIsOwner);
fCurrentBasket = nullptr;
fBaskets[fReadBasket] = nullptr;
} else {
// This is the only copy, we can't return it as is to the user, just make a copy.
if (user_buf.BufferSize() < buf->BufferSize()) {
user_buf.AutoExpand(buf->BufferSize());
}
memcpy(user_buf.Buffer(), buf->Buffer(), buf->BufferSize());
}
}

Int_t bufbegin = basket->GetKeylen();
buf->SetBufferOffset(bufbegin);
user_buf.SetBufferOffset(bufbegin);

Int_t N = ((fNextBasketEntry < 0) ? fEntryNumber : fNextBasketEntry) - first;
//printf("Requesting %d events; fNextBasketEntry=%lld; first=%lld.\n", N, fNextBasketEntry, first);
if (R__unlikely(!leaf->ReadBasketFast(*buf, N))) {
if (R__unlikely(!leaf->ReadBasketFast(user_buf, N))) {
Error("GetBulkEntries", "Leaf failed to read.\n");
return -1;
}
user_buf.SetBufferOffset(bufbegin);

fCurrentBasket = nullptr;
fBaskets[fReadBasket] = nullptr;
R__ASSERT(fExtraBasket == nullptr && "fExtraBasket should have been set to nullptr by GetFreshBasket");
fExtraBasket = basket;
basket->DisownBuffer();
if (fCurrentBasket == nullptr) {
R__ASSERT(fExtraBasket == nullptr && "fExtraBasket should have been set to nullptr by GetFreshBasket");
fExtraBasket = basket;
basket->DisownBuffer();
}

return N;
}
Expand All @@ -1498,7 +1522,7 @@ Int_t TBranch::GetEntriesSerialized(Long64_t entry, TBuffer &user_buf, TBuffer *
TBasket *basket = nullptr;
Long64_t first;
Int_t result = GetBasketAndFirst(basket, first, &user_buf);
if (R__unlikely(result <= 0)) { return -1; }
if (R__unlikely(result < 0)) { return -1; }
// Only support reading from full clusters.
if (R__unlikely(entry != first)) {
Error("GetEntriesSerialized", "Failed to read from full cluster; first entry is %lld; requested entry is %lld.\n", first, entry);
Expand All @@ -1519,13 +1543,32 @@ Int_t TBranch::GetEntriesSerialized(Long64_t entry, TBuffer &user_buf, TBuffer *
return -1;
}

if (&user_buf != buf) {
// The basket was already in memory and might (and might not) be backed by persistent
// storage.
R__ASSERT(result == fReadBasket);
if (fBasketSeek[fReadBasket]) {
// It is backed, so we can be destructive
user_buf.SetBuffer(buf->Buffer(), buf->BufferSize());
buf->ResetBit(TBufferIO::kIsOwner);
fCurrentBasket = nullptr;
fBaskets[fReadBasket] = nullptr;
} else {
// This is the only copy, we can't return it as is to the user, just make a copy.
if (user_buf.BufferSize() < buf->BufferSize()) {
user_buf.AutoExpand(buf->BufferSize());
}
memcpy(user_buf.Buffer(), buf->Buffer(), buf->BufferSize());
}
}

Int_t bufbegin = basket->GetKeylen();
buf->SetBufferOffset(bufbegin);
user_buf.SetBufferOffset(bufbegin);

Int_t N = ((fNextBasketEntry < 0) ? fEntryNumber : fNextBasketEntry) - first;
//Info("GetEntriesSerialized", "Requesting %d events; fNextBasketEntry=%lld; first=%lld.\n", N, fNextBasketEntry, first);

if (R__unlikely(!leaf->ReadBasketSerialized(*buf, N))) {
if (R__unlikely(!leaf->ReadBasketSerialized(user_buf, N))) {
Error("GetEntriesSerialized", "Leaf failed to read.\n");
return -1;
}
Expand Down Expand Up @@ -1587,7 +1630,7 @@ Int_t TBranch::GetEntry(Long64_t entry, Int_t getall)
Long64_t first;

Int_t result = GetBasketAndFirst(basket, first, nullptr);
if (R__unlikely(result <= 0)) { return result; }
if (R__unlikely(result < 0)) { return result + 1; }

basket->PrepareBasket(entry);
TBuffer* buf = basket->GetBufferRef();
Expand Down

0 comments on commit fe9a8a6

Please sign in to comment.