fixup! rework index builder

Author: aalkin

Framework/AnalysisSupport/src/AODReaderHelpers.cxx

@@ -72,7 +72,7 @@ struct Buildable {
       outputSchema,
       origin,
       description,
-      version};
+      version, nullptr};
   }
 
 };
@@ -91,7 +91,7 @@ AlgorithmSpec AODReaderHelpers::indexBuilderCallback(ConfigContext const& ctx)
     for (auto& b : buildables) {
       builders.push_back(b.createBuilder());
     }
-    return [builders](ProcessingContext& pc) {
+    return [builders](ProcessingContext& pc) mutable {
       auto outputs = pc.outputs();
       for (auto& builder : builders) {
         outputs.adopt(Output{builder.origin, builder.description, builder.version}, builder.materialize(pc));

Framework/Core/include/Framework/AnalysisHelpers.h

@@ -60,7 +60,11 @@ struct IndexRecord {
 };
 
 struct IndexBuilder {
-  static std::shared_ptr<arrow::Table> materialize(std::vector<std::shared_ptr<arrow::Table>>&& tables, std::vector<soa::IndexRecord> const& records, std::shared_ptr<arrow::Schema> const& schema, bool exclusive);
+  static std::vector<framework::IndexColumnBuilderNG> makeBuilders(std::vector<std::shared_ptr<arrow::Table>>&& tables, std::vector<soa::IndexRecord> const& records);
+  static void resetBuilders(std::vector<framework::IndexColumnBuilderNG>& builders, std::vector<std::shared_ptr<arrow::Table>>&& tables);
+
+  // static std::shared_ptr<arrow::Table> materialize(std::vector<std::shared_ptr<arrow::Table>>&& tables, std::vector<soa::IndexRecord> const& records, std::shared_ptr<arrow::Schema> const& schema, bool exclusive);
+  static std::shared_ptr<arrow::Table> materializeNG(std::vector<framework::IndexColumnBuilderNG>& builders, std::vector<std::shared_ptr<arrow::Table>>&& tables, std::vector<soa::IndexRecord> const& records, std::shared_ptr<arrow::Schema> const& schema, bool exclusive);
 };
 } // namespace o2::soa
 
@@ -161,7 +165,9 @@ struct Builder {
   header::DataDescription description;
   header::DataHeader::SubSpecificationType version;
 
-  std::shared_ptr<arrow::Table> materialize(ProcessingContext& pc) const;
+  std::shared_ptr<std::vector<framework::IndexColumnBuilderNG>> builders = nullptr;
+
+  std::shared_ptr<arrow::Table> materialize(ProcessingContext& pc);
 };
 }  // namespace o2::framework
 
@@ -722,6 +728,8 @@ struct Builds : decltype(transformBase<T>()) {
 
   std::vector<soa::IndexRecord> map = soa::getIndexMapping<metadata>();
 
+  std::vector<framework::IndexColumnBuilderNG> builders;
+
   T* operator->()
   {
     return table.get();
@@ -744,7 +752,7 @@ struct Builds : decltype(transformBase<T>()) {
 
   auto build(std::vector<std::shared_ptr<arrow::Table>>&& tables)
   {
-    this->table = std::make_shared<T>(soa::IndexBuilder::materialize(std::forward<std::vector<std::shared_ptr<arrow::Table>>>(tables), map, outputSchema, metadata::exclusive));
+    this->table = std::make_shared<T>(soa::IndexBuilder::materializeNG(builders, std::forward<std::vector<std::shared_ptr<arrow::Table>>>(tables), map, outputSchema, metadata::exclusive));
     return (this->table != nullptr);
   }
 };

Framework/Core/include/Framework/IndexBuilderHelpers.h

@@ -14,7 +14,6 @@
 #include <arrow/chunked_array.h>
 #include <arrow/builder.h>
 #include <arrow/memory_pool.h>
-#include <string>
 #include <memory>
 
 namespace o2::soa
@@ -35,6 +34,7 @@ void cannotCreateIndexBuilder();
 
 struct ChunkedArrayIterator {
   ChunkedArrayIterator(std::shared_ptr<arrow::ChunkedArray> source);
+  void reset(std::shared_ptr<arrow::ChunkedArray>& source);
 
   std::shared_ptr<arrow::ChunkedArray> mSource = nullptr;
   size_t mPosition = 0;
@@ -44,6 +44,7 @@ struct ChunkedArrayIterator {
   int const* mCurrent = nullptr;
   int const* mLast = nullptr;
   size_t mFirstIndex = 0;
+  size_t mSourceSize = 0;
 
   std::shared_ptr<arrow::Int32Array> getCurrentArray();
   void nextChunk();
@@ -53,160 +54,170 @@ struct ChunkedArrayIterator {
 
 struct SelfBuilder {
   std::unique_ptr<arrow::ArrayBuilder> mBuilder = nullptr;
+  std::unique_ptr<framework::ChunkedArrayIterator> keyIndex = nullptr;
   SelfBuilder(arrow::MemoryPool* pool);
+  void reset(std::shared_ptr<arrow::ChunkedArray>);
+
+  inline bool find(int) const
+  {
+    return true;
+  }
+  void fill(int idx);
+  std::shared_ptr<arrow::ChunkedArray> result() const;
 };
 
-struct SingleBuilder {
-  ChunkedArrayIterator arrayIterator;
+struct SingleBuilder : public ChunkedArrayIterator {
   std::unique_ptr<arrow::ArrayBuilder> mBuilder = nullptr;
   SingleBuilder(std::shared_ptr<arrow::ChunkedArray> source, arrow::MemoryPool* pool);
+  void reset(std::shared_ptr<arrow::ChunkedArray> source);
+
+  bool find(int idx);
+  void fill(int idx);
+  std::shared_ptr<arrow::ChunkedArray> result() const;
 };
 
-struct SliceBuilder {
-  ChunkedArrayIterator arrayIterator;
+struct SliceBuilder : public ChunkedArrayIterator {
   arrow::ArrayBuilder* mValueBuilder = nullptr;
   std::unique_ptr<arrow::ArrayBuilder> mListBuilder = nullptr;
   std::shared_ptr<arrow::NumericArray<arrow::Int32Type>> mValues = nullptr;
   std::shared_ptr<arrow::NumericArray<arrow::Int64Type>> mCounts = nullptr;
+  int mValuePos = 0;
   SliceBuilder(std::shared_ptr<arrow::ChunkedArray> source, arrow::MemoryPool* pool);
+  void reset(std::shared_ptr<arrow::ChunkedArray> source);
+
+  bool find(int idx);
+  void fill(int idx);
+  std::shared_ptr<arrow::ChunkedArray> result() const;
 
   arrow::Status preSlice();
 };
 
-struct ArrayBuilder {
-  ChunkedArrayIterator arrayIterator;
+struct ArrayBuilder : public ChunkedArrayIterator {
   arrow::ArrayBuilder* mValueBuilder = nullptr;
   std::vector<int> mValues;
   std::vector<std::vector<int>> mIndices;
   std::unique_ptr<arrow::ArrayBuilder> mListBuilder = nullptr;
   ArrayBuilder(std::shared_ptr<arrow::ChunkedArray> source, arrow::MemoryPool* pool);
+  void reset(std::shared_ptr<arrow::ChunkedArray> source);
+
+  bool find(int idx);
+  void fill(int idx);
+  std::shared_ptr<arrow::ChunkedArray> result() const;
 
   arrow::Status preFind();
 };
 
 struct IndexColumnBuilderNG {
   std::variant<std::monostate, SelfBuilder, SingleBuilder, SliceBuilder, ArrayBuilder> builder;
-
-  IndexColumnBuilderNG(soa::IndexKind kind, arrow::MemoryPool* pool, std::shared_ptr<arrow::ChunkedArray> source = nullptr)
-  {
-    switch (kind) {
-      case soa::IndexKind::IdxSelf:
-        builder = SelfBuilder{pool};
-        break;
-      case soa::IndexKind::IdxSingle:
-        builder = SingleBuilder{source, pool};
-        break;
-      case soa::IndexKind::IdxSlice:
-        builder = SliceBuilder{source, pool};
-        break;
-      case soa::IndexKind::IdxArray:
-        builder = ArrayBuilder{source, pool};
-        break;
-      default:
-        cannotCreateIndexBuilder();
-    }
-  }
-};
-
-struct SelfIndexColumnBuilder {
-  SelfIndexColumnBuilder(const char* name, arrow::MemoryPool* pool);
-  virtual ~SelfIndexColumnBuilder() = default;
-
-  inline std::shared_ptr<arrow::ChunkedArray> result() const
-  {
-    std::shared_ptr<arrow::Array> array;
-    auto status = static_cast<arrow::Int32Builder*>(mBuilder.get())->Finish(&array);
-    if (!status.ok()) {
-      cannotBuildAnArray();
-    }
-
-    return std::make_shared<arrow::ChunkedArray>(array);
-  }
-
-  inline bool find(int)
-  {
-    return true;
-  }
-
-  inline void fill(int idx)
-  {
-    (void)static_cast<arrow::Int32Builder*>(mBuilder.get())->Append(idx);
-  }
-
-  std::string mColumnName;
-  std::unique_ptr<arrow::ArrayBuilder> mBuilder = nullptr;
-};
-
-class IndexColumnBuilder : public SelfIndexColumnBuilder, public ChunkedArrayIterator
-{
- public:
-  IndexColumnBuilder(std::shared_ptr<arrow::ChunkedArray> source, const char* name, int listSize, arrow::MemoryPool* pool);
-  ~IndexColumnBuilder() override = default;
-
-  inline std::shared_ptr<arrow::ChunkedArray> result() const
-  {
-    if (mListSize == -1) {
-      return resultMulti();
-    } else if (mListSize == 2) {
-      return resultSlice();
-    } else {
-      return resultSingle();
-    }
-  }
-
-  inline bool find(int idx)
-  {
-    if (mListSize == -1) {
-      return findMulti(idx);
-    } else if (mListSize == 2) {
-      return findSlice(idx);
-    } else {
-      return findSingle(idx);
-    }
-  }
-
-  inline void fill(int idx)
-  {
-    ++mResultSize;
-    if (mListSize == -1) {
-      fillMulti(idx);
-    } else if (mListSize == 2) {
-      fillSlice(idx);
-    } else {
-      fillSingle(idx);
-    }
-  }
-
- private:
-  arrow::Status preSlice();
-  arrow::Status preFind();
-
-  bool findSingle(int idx);
-  bool findSlice(int idx);
-  bool findMulti(int idx);
-
-  void fillSingle(int idx);
-  void fillSlice(int idx);
-  void fillMulti(int idx);
-
-  std::shared_ptr<arrow::ChunkedArray> resultSingle() const;
-  std::shared_ptr<arrow::ChunkedArray> resultSlice() const;
-  std::shared_ptr<arrow::ChunkedArray> resultMulti() const;
-
-  int mListSize = 1;
-  arrow::ArrayBuilder* mValueBuilder = nullptr;
-  std::unique_ptr<arrow::ArrayBuilder> mListBuilder = nullptr;
-
-  size_t mSourceSize = 0;
   size_t mResultSize = 0;
+  int mColumnPos = -1;
+  IndexColumnBuilderNG(soa::IndexKind kind, int pos, arrow::MemoryPool* pool, std::shared_ptr<arrow::ChunkedArray> source = nullptr);
+  void reset(std::shared_ptr<arrow::ChunkedArray> source = nullptr);
 
-  std::shared_ptr<arrow::NumericArray<arrow::Int32Type>> mValuesArrow = nullptr;
-  std::shared_ptr<arrow::NumericArray<arrow::Int64Type>> mCounts = nullptr;
-  std::vector<int> mValues;
-  std::vector<std::vector<int>> mIndices;
-  int mFillOffset = 0;
-  int mValuePos = 0;
+  bool find(int idx);
+  void fill(int idx);
+  std::shared_ptr<arrow::ChunkedArray> result() const;
 };
+
+// struct SelfIndexColumnBuilder {
+//   SelfIndexColumnBuilder(const char* name, arrow::MemoryPool* pool);
+//   virtual ~SelfIndexColumnBuilder() = default;
+
+//   inline std::shared_ptr<arrow::ChunkedArray> result() const
+//   {
+//     std::shared_ptr<arrow::Array> array;
+//     auto status = static_cast<arrow::Int32Builder*>(mBuilder.get())->Finish(&array);
+//     if (!status.ok()) {
+//       cannotBuildAnArray();
+//     }
+
+//     return std::make_shared<arrow::ChunkedArray>(array);
+//   }
+
+//   inline bool find(int)
+//   {
+//     return true;
+//   }
+
+//   inline void fill(int idx)
+//   {
+//     (void)static_cast<arrow::Int32Builder*>(mBuilder.get())->Append(idx);
+//   }
+
+//   std::string mColumnName;
+//   std::unique_ptr<arrow::ArrayBuilder> mBuilder = nullptr;
+// };
+
+// class IndexColumnBuilder : public SelfIndexColumnBuilder, public ChunkedArrayIterator
+// {
+//  public:
+//   IndexColumnBuilder(std::shared_ptr<arrow::ChunkedArray> source, const char* name, int listSize, arrow::MemoryPool* pool);
+//   ~IndexColumnBuilder() override = default;
+
+//   inline std::shared_ptr<arrow::ChunkedArray> result() const
+//   {
+//     if (mListSize == -1) {
+//       return resultMulti();
+//     } else if (mListSize == 2) {
+//       return resultSlice();
+//     } else {
+//       return resultSingle();
+//     }
+//   }
+
+//   inline bool find(int idx)
+//   {
+//     if (mListSize == -1) {
+//       return findMulti(idx);
+//     } else if (mListSize == 2) {
+//       return findSlice(idx);
+//     } else {
+//       return findSingle(idx);
+//     }
+//   }
+
+//   inline void fill(int idx)
+//   {
+//     ++mResultSize;
+//     if (mListSize == -1) {
+//       fillMulti(idx);
+//     } else if (mListSize == 2) {
+//       fillSlice(idx);
+//     } else {
+//       fillSingle(idx);
+//     }
+//   }
+
+//  private:
+//   arrow::Status preSlice();
+//   arrow::Status preFind();
+
+//   bool findSingle(int idx);
+//   bool findSlice(int idx);
+//   bool findMulti(int idx);
+
+//   void fillSingle(int idx);
+//   void fillSlice(int idx);
+//   void fillMulti(int idx);
+
+//   std::shared_ptr<arrow::ChunkedArray> resultSingle() const;
+//   std::shared_ptr<arrow::ChunkedArray> resultSlice() const;
+//   std::shared_ptr<arrow::ChunkedArray> resultMulti() const;
+
+//   int mListSize = 1;
+//   arrow::ArrayBuilder* mValueBuilder = nullptr;
+//   std::unique_ptr<arrow::ArrayBuilder> mListBuilder = nullptr;
+
+//   size_t mSourceSize = 0;
+//   size_t mResultSize = 0;
+
+//   std::shared_ptr<arrow::NumericArray<arrow::Int32Type>> mValuesArrow = nullptr;
+//   std::shared_ptr<arrow::NumericArray<arrow::Int64Type>> mCounts = nullptr;
+//   std::vector<int> mValues;
+//   std::vector<std::vector<int>> mIndices;
+//   int mFillOffset = 0;
+//   int mValuePos = 0;
+// };
 } // namespace o2::framework
 
 #endif // O2_FRAMEWORK_INDEXBUILDERHELPERS_H_