Merge f0594c0 into sapling-pr-archive-ktf

Author: ktf

Framework/Core/include/Framework/TableTreeHelpers.h

@@ -45,29 +45,6 @@ struct ROOTTypeInfo {
 auto arrowTypeFromROOT(EDataType type, int size);
 auto basicROOTTypeFromArrow(arrow::Type::type id);
 
-class BranchToColumn
-{
- public:
-  BranchToColumn(TBranch* branch, bool VLA, std::string name, EDataType type, int listSize, arrow::MemoryPool* pool);
-  //  BranchToColumn(TBranch* branch, TBranch* sizeBranch, std::string name, EDataType type, arrow::MemoryPool* pool);
-  ~BranchToColumn() = default;
-  TBranch* branch();
-
-  std::pair<std::shared_ptr<arrow::ChunkedArray>, std::shared_ptr<arrow::Field>> read(TBuffer* buffer);
-
- private:
-  TBranch* mBranch = nullptr;
-  bool mVLA = false;
-  std::string mColumnName;
-  EDataType mType;
-  std::shared_ptr<arrow::DataType> mArrowType;
-  arrow::ArrayBuilder* mValueBuilder = nullptr;
-  std::unique_ptr<arrow::ArrayBuilder> mListBuilder = nullptr;
-  int mListSize = 1;
-  std::unique_ptr<arrow::ArrayBuilder> mBuilder = nullptr;
-  arrow::MemoryPool* mPool = nullptr;
-};
-
 class ColumnToBranch
 {
  public:

Framework/Core/src/TableTreeHelpers.cxx

@@ -102,166 +102,6 @@ auto basicROOTTypeFromArrow(arrow::Type::type id)
   }
 }
 
-TBranch* BranchToColumn::branch()
-{
-  return mBranch;
-}
-
-BranchToColumn::BranchToColumn(TBranch* branch, bool VLA, std::string name, EDataType type, int listSize, arrow::MemoryPool* pool)
-  : mBranch{branch},
-    mVLA{VLA},
-    mColumnName{std::move(name)},
-    mType{type},
-    mArrowType{arrowTypeFromROOT(type, listSize)},
-    mListSize{listSize},
-    mPool{pool}
-
-{
-  if (mType == EDataType::kBool_t) {
-    if (mListSize > 1) {
-      auto status = arrow::MakeBuilder(mPool, mArrowType->field(0)->type(), &mBuilder);
-      if (!status.ok()) {
-        throw runtime_error("Cannot create value builder");
-      }
-      mListBuilder = std::make_unique<arrow::FixedSizeListBuilder>(mPool, std::move(mBuilder), mListSize);
-      mValueBuilder = static_cast<arrow::FixedSizeListBuilder*>(mListBuilder.get())->value_builder();
-    } else {
-      auto status = arrow::MakeBuilder(mPool, mArrowType, &mBuilder);
-      if (!status.ok()) {
-        throw runtime_error("Cannot create builder");
-      }
-      mValueBuilder = mBuilder.get();
-    }
-  }
-}
-
-std::pair<std::shared_ptr<arrow::ChunkedArray>, std::shared_ptr<arrow::Field>> BranchToColumn::read(TBuffer* buffer)
-{
-  O2_SIGNPOST_ID_FROM_POINTER(sid, tabletree_helpers, buffer);
-  auto totalEntries = mBranch->GetEntries();
-  arrow::Status status;
-  int readEntries = 0;
-  buffer->Reset();
-  std::shared_ptr<arrow::Array> array;
-
-  if (mType == EDataType::kBool_t) {
-    // boolean array special case: we need to use builder to create the bitmap
-    status = mValueBuilder->Reserve(totalEntries * mListSize);
-    if (mListSize > 1) {
-      status &= mListBuilder->Reserve(totalEntries);
-    }
-    if (!status.ok()) {
-      throw runtime_error("Failed to reserve memory for array builder");
-    }
-    while (readEntries < totalEntries) {
-      auto readLast = mBranch->GetBulkRead().GetBulkEntries(readEntries, *buffer);
-      readEntries += readLast;
-      status &= static_cast<arrow::BooleanBuilder*>(mValueBuilder)->AppendValues(reinterpret_cast<uint8_t const*>(buffer->GetCurrent()), readLast * mListSize);
-    }
-    if (mListSize > 1) {
-      status &= static_cast<arrow::FixedSizeListBuilder*>(mListBuilder.get())->AppendValues(readEntries);
-    }
-    if (!status.ok()) {
-      throw runtime_error("Failed to append values to array");
-    }
-    if (mListSize > 1) {
-      status &= mListBuilder->Finish(&array);
-    } else {
-      status &= mValueBuilder->Finish(&array);
-    }
-    if (!status.ok()) {
-      throw runtime_error("Failed to create array");
-    }
-  } else {
-    // other types: use serialized read to build arrays directly
-    size_t branchSize = mBranch->GetTotBytes();
-    auto&& result = arrow::AllocateResizableBuffer(mBranch->GetTotBytes(), mPool);
-    O2_SIGNPOST_EVENT_EMIT(tabletree_helpers, sid, "BranchToColumn", "Allocating %ld bytes for %{public}s", branchSize, mBranch->GetName());
-    if (!result.ok()) {
-      throw runtime_error("Cannot allocate values buffer");
-    }
-    std::shared_ptr<arrow::Buffer> arrowValuesBuffer = std::move(result).ValueUnsafe();
-    auto ptr = arrowValuesBuffer->mutable_data();
-    if (ptr == nullptr) {
-      throw runtime_error("Invalid buffer");
-    }
-
-    auto typeSize = TDataType::GetDataType(mType)->Size();
-    std::unique_ptr<TBufferFile> offsetBuffer = nullptr;
-
-    uint32_t offset = 0;
-    int count = 0;
-    std::shared_ptr<arrow::Buffer> arrowOffsetBuffer;
-    gsl::span<int> offsets;
-    int size = 0;
-    uint32_t totalSize = 0;
-    TBranch* mSizeBranch = nullptr;
-    if (mVLA) {
-      mSizeBranch = mBranch->GetTree()->GetBranch((std::string{mBranch->GetName()} + TableTreeHelpers::sizeBranchSuffix).c_str());
-      offsetBuffer = std::make_unique<TBufferFile>(TBuffer::EMode::kWrite, 4 * 1024 * 1024);
-      result = arrow::AllocateResizableBuffer((totalEntries + 1) * (int64_t)sizeof(int), mPool);
-      if (!result.ok()) {
-        throw runtime_error("Cannot allocate offset buffer");
-      }
-      arrowOffsetBuffer = std::move(result).ValueUnsafe();
-      unsigned char* ptrOffset = arrowOffsetBuffer->mutable_data();
-      auto* tPtrOffset = reinterpret_cast<int*>(ptrOffset);
-      offsets = gsl::span<int>{tPtrOffset, tPtrOffset + totalEntries + 1};
-
-      // read sizes first
-      while (readEntries < totalEntries) {
-        auto readLast = mSizeBranch->GetBulkRead().GetEntriesSerialized(readEntries, *offsetBuffer);
-        readEntries += readLast;
-        for (auto i = 0; i < readLast; ++i) {
-          offsets[count++] = (int)offset;
-          offset += swap32_(reinterpret_cast<uint32_t*>(offsetBuffer->GetCurrent())[i]);
-        }
-      }
-      offsets[count] = (int)offset;
-      totalSize = offset;
-      readEntries = 0;
-    }
-
-    while (readEntries < totalEntries) {
-      auto readLast = mBranch->GetBulkRead().GetEntriesSerialized(readEntries, *buffer);
-      if (mVLA) {
-        size = offsets[readEntries + readLast] - offsets[readEntries];
-      } else {
-        size = readLast * mListSize;
-      }
-      readEntries += readLast;
-      swapCopy(ptr, buffer->GetCurrent(), size, typeSize);
-      ptr += (ptrdiff_t)(size * typeSize);
-    }
-    if (!mVLA) {
-      totalSize = readEntries * mListSize;
-    }
-    std::shared_ptr<arrow::PrimitiveArray> varray;
-    switch (mListSize) {
-      case -1:
-        varray = std::make_shared<arrow::PrimitiveArray>(mArrowType->field(0)->type(), totalSize, arrowValuesBuffer);
-        array = std::make_shared<arrow::ListArray>(mArrowType, readEntries, arrowOffsetBuffer, varray);
-        break;
-      case 1:
-        array = std::make_shared<arrow::PrimitiveArray>(mArrowType, readEntries, arrowValuesBuffer);
-        break;
-      default:
-        varray = std::make_shared<arrow::PrimitiveArray>(mArrowType->field(0)->type(), totalSize, arrowValuesBuffer);
-        array = std::make_shared<arrow::FixedSizeListArray>(mArrowType, readEntries, varray);
-    }
-  }
-
-  auto fullArray = std::make_shared<arrow::ChunkedArray>(array);
-  auto field = std::make_shared<arrow::Field>(mBranch->GetName(), mArrowType);
-
-  mBranch->SetStatus(false);
-  mBranch->DropBaskets("all");
-  mBranch->Reset();
-  mBranch->GetTransientBuffer(0)->Expand(0);
-
-  return std::make_pair(fullArray, field);
-}
-
 ColumnToBranch::ColumnToBranch(TTree* tree, std::shared_ptr<arrow::ChunkedArray> const& column, std::shared_ptr<arrow::Field> const& field)
   : mBranchName{field->name()},
     mColumn{column.get()},