DPL Analysis: avoid Arrow's Slice API

Framework/Core/include/Framework/AnalysisHelpers.h

@@ -147,7 +147,7 @@ auto spawner(framework::pack<C...>, std::vector<std::shared_ptr<arrow::Table>>&&
   if (fullTable->num_rows() == 0) {
     return makeEmptyTable(name, framework::pack<C...>{});
   }
-  return spawnerHelper(fullTable, schema, sizeof...(C), projectors, name, projector);
+  return spawnerHelper(fullTable.tablePtr, schema, sizeof...(C), projectors, name, projector);
 }
 
 std::string serializeProjectors(std::vector<framework::expressions::Projector>& projectors);
@@ -950,7 +950,7 @@ auto getTableFromFilter(soa::is_filtered_table auto const& table, soa::Selection
 
 auto getTableFromFilter(soa::is_not_filtered_table auto const& table, soa::SelectionVector&& selection)
 {
-  return std::make_unique<o2::soa::Filtered<std::decay_t<decltype(table)>>>(std::vector{table.asArrowTable()}, std::forward<soa::SelectionVector>(selection));
+  return std::make_unique<o2::soa::Filtered<std::decay_t<decltype(table)>>>(std::vector{table.asArrowTableRef()}, std::forward<soa::SelectionVector>(selection));
 }
 
 void initializePartitionCaches(std::set<uint32_t> const& hashes, std::shared_ptr<arrow::Schema> const& schema, expressions::Filter const& filter, gandiva::NodePtr& tree, gandiva::FilterPtr& gfilter);
@@ -982,7 +982,7 @@ struct Partition {
 
   void bindTable(T const& table)
   {
-    intializeCaches(T::table_t::hashes(), table.asArrowTable()->schema());
+    intializeCaches(T::table_t::hashes(), table.asArrowTableRef()->schema());
     if (dataframeChanged) {
       mFiltered = getTableFromFilter(table, soa::selectionToVector(framework::expressions::createSelection(table.asArrowTable(), gfilter)));
       dataframeChanged = false;
@@ -1086,7 +1086,7 @@ auto Extend(T const& table)
   static std::array<framework::expressions::Projector, sizeof...(Cs)> projectors{{std::move(Cs::Projector())...}};
   static std::shared_ptr<gandiva::Projector> projector = nullptr;
   static auto schema = std::make_shared<arrow::Schema>(o2::soa::createFieldsFromColumns(framework::pack<Cs...>{}));
-  return output_t{{o2::framework::spawner(framework::pack<Cs...>{}, {table.asArrowTable()}, "dynamicExtension", projectors.data(), projector, schema), table.asArrowTable()}, 0};
+  return output_t{{o2::framework::spawner(framework::pack<Cs...>{}, {table.asArrowTable()}, "dynamicExtension", projectors.data(), projector, schema), table.asArrowTable()}};
 }
 
 /// Template function to attach dynamic columns on-the-fly (e.g. inside
@@ -1095,7 +1095,7 @@ template <soa::is_table T, soa::is_dynamic_column... Cs>
 auto Attach(T const& table)
 {
   using output_t = Join<T, o2::soa::Table<o2::aod::Hash<"JOIN"_h>, o2::aod::Hash<"JOIN/0"_h>, o2::aod::Hash<"JOIN"_h>, Cs...>>;
-  return output_t{{table.asArrowTable()}, table.offset()};
+  return output_t{{table.asArrowTableRef()}};
 }
 } // namespace o2::soa
 

Framework/Core/include/Framework/AnalysisManagers.h

@@ -319,12 +319,12 @@ bool prepareOutput(ProcessingContext& context, T& spawns)
   }
   using D = o2::aod::Hash<metadata::extension_table_t::ref.desc_hash>;
 
-  spawns.extension = std::make_shared<typename T::extension_t>(o2::framework::spawner<D>(originalTable,
+  spawns.extension = std::make_shared<typename T::extension_t>(o2::framework::spawner<D>(originalTable.tablePtr,
                                                                                          o2::aod::label<metadata::extension_table_t::ref>(),
                                                                                          spawns.projectors.data(),
                                                                                          spawns.projector,
                                                                                          spawns.schema));
-  spawns.table = std::make_shared<typename T::spawnable_t::table_t>(soa::ArrowHelpers::joinTables({spawns.extension->asArrowTable(), originalTable}, std::span{T::spawnable_t::table_t::originalLabels}));
+  spawns.table = std::make_shared<typename T::spawnable_t::table_t>(soa::ArrowHelpers::joinTables({spawns.extension->asArrowTableRef(), originalTable}, std::span{T::spawnable_t::table_t::originalLabels}));
   return true;
 }
 
@@ -348,12 +348,12 @@ bool prepareOutput(ProcessingContext& context, T& defines)
   }
   using D = o2::aod::Hash<metadata::extension_table_t::ref.desc_hash>;
 
-  defines.extension = std::make_shared<typename T::extension_t>(o2::framework::spawner<D>(originalTable,
+  defines.extension = std::make_shared<typename T::extension_t>(o2::framework::spawner<D>(originalTable.tablePtr,
                                                                                           o2::aod::label<metadata::extension_table_t::ref>(),
                                                                                           defines.projectors.data(),
                                                                                           defines.projector,
                                                                                           defines.schema));
-  defines.table = std::make_shared<typename T::spawnable_t::table_t>(soa::ArrowHelpers::joinTables({defines.extension->asArrowTable(), originalTable}, std::span{T::spawnable_t::table_t::originalLabels}));
+  defines.table = std::make_shared<typename T::spawnable_t::table_t>(soa::ArrowHelpers::joinTables({defines.extension->asArrowTableRef(), originalTable}, std::span{T::spawnable_t::table_t::originalLabels}));
   return true;
 }
 
@@ -380,12 +380,12 @@ bool prepareDelayedOutput(ProcessingContext& context, T& defines)
   }
   using D = o2::aod::Hash<metadata::extension_table_t::ref.desc_hash>;
 
-  defines.extension = std::make_shared<typename T::extension_t>(o2::framework::spawner<D>(originalTable,
+  defines.extension = std::make_shared<typename T::extension_t>(o2::framework::spawner<D>(originalTable.tablePtr,
                                                                                           o2::aod::label<metadata::extension_table_t::ref>(),
                                                                                           defines.projectors.data(),
                                                                                           defines.projector,
                                                                                           defines.schema));
-  defines.table = std::make_shared<typename T::spawnable_t::table_t>(soa::ArrowHelpers::joinTables({defines.extension->asArrowTable(), originalTable}, std::span{T::spawnable_t::table_t::originalLabels}));
+  defines.table = std::make_shared<typename T::spawnable_t::table_t>(soa::ArrowHelpers::joinTables({defines.extension->asArrowTableRef(), originalTable}, std::span{T::spawnable_t::table_t::originalLabels}));
   return true;
 }
 

Framework/Core/include/Framework/AnalysisTask.h

@@ -226,7 +226,7 @@ struct AnalysisDataProcessorBuilder {
   template <std::ranges::input_range R>
   static auto extractTablesFromRecord(InputRecord& record, R matchers)
   {
-    std::vector<std::shared_ptr<arrow::Table>> tables;
+    std::vector<soa::ArrowTableRef> tables;
     std::ranges::transform(matchers, std::back_inserter(tables), [&record](auto const& m) {
       return record.get<TableConsumer>(m.second)->asArrowTable();
     });
@@ -248,8 +248,8 @@ struct AnalysisDataProcessorBuilder {
   template <soa::is_filtered T, std::ranges::input_range R>
   static auto extractFilteredFromRecord(InputRecord& record, R matchers, ExpressionInfo& info)
   {
-    std::shared_ptr<arrow::Table> table = soa::ArrowHelpers::joinTables(extractTablesFromRecord(record, matchers));
-    expressions::updateFilterInfo(info, table);
+    auto table = soa::ArrowHelpers::joinTables(extractTablesFromRecord(record, matchers));
+    expressions::updateFilterInfo(info, table.tablePtr);
     if constexpr (!o2::soa::is_smallgroups<std::decay_t<T>>) {
       if (info.selection == nullptr) {
         soa::missingFilterDeclaration(info.processHash, info.argumentIndex);

Framework/Core/include/Framework/ArrowTypes.h

@@ -11,12 +11,55 @@
 
 #ifndef O2_FRAMEWORK_ARROWTYPES_H
 #define O2_FRAMEWORK_ARROWTYPES_H
+#include <arrow/table.h>
 #include "Framework/Traits.h"
 #include "arrow/type_fwd.h"
 #include <span>
 
 namespace o2::soa
 {
+struct ArrowRange {
+  uint64_t offset;
+  int64_t size;
+
+  bool operator!=(ArrowRange const& other) const
+  {
+    return (offset != other.offset) && (size != other.size);
+  }
+};
+
+struct ArrowTableRef {
+  std::shared_ptr<arrow::Table> tablePtr = nullptr;
+  ArrowRange range{0, 0};
+
+  ArrowTableRef() = default;
+  ArrowTableRef(std::shared_ptr<arrow::Table> table)
+    : tablePtr{table},
+      range{0, table->num_rows()}
+  {
+  }
+  ArrowTableRef(std::shared_ptr<arrow::Table> table, ArrowRange range_)
+    : tablePtr{table},
+      range{range_}
+  {
+  }
+
+  ArrowTableRef makeEmpty() const
+  {
+    return {tablePtr, {0, 0}};
+  }
+
+  ArrowTableRef slice(ArrowRange newRange) const
+  {
+    return {tablePtr, newRange};
+  }
+
+  std::shared_ptr<arrow::Table> const& operator->() const
+  {
+    return tablePtr;
+  }
+};
+
 template <typename T>
 struct arrow_array_for {
 };

Framework/Core/include/Framework/GroupSlicer.h

@@ -194,7 +194,7 @@ struct GroupSlicer {
           }
         }
       }
-      std::decay_t<A1> typedTable{{originalTable.asArrowTable()}, std::move(s)};
+      std::decay_t<A1> typedTable{{originalTable.asArrowTableRef()}, std::move(s)};
       typedTable.bindInternalIndicesTo(&originalTable);
       return typedTable;
     }
@@ -218,16 +218,7 @@ struct GroupSlicer {
       auto oc = sliceInfos[index].getSliceFor(pos);
       uint64_t offset = oc.first;
       auto count = oc.second;
-      if (count == 0) {
-        // Empty group: avoid slicing every column only to discard it. Cache one
-        // empty (0-row) table per associated table and reuse it. This is the
-        // common case for sparse grouping (e.g. collisions with no candidates).
-        if (!emptyTables[index]) {
-          emptyTables[index] = originalTable.asArrowTable()->Slice(0, 0);
-        }
-        return std::decay_t<A1>{{emptyTables[index]}, soa::SelectionVector{}};
-      }
-      auto groupedElementsTable = originalTable.asArrowTable()->Slice(offset, count);
+      auto groupedElementsTable = originalTable.asArrowTableRef().slice({offset, count});
 
       // for each grouping element we need to slice the selection vector
       auto start_iterator = std::lower_bound(starts[index], selections[index]->end(), offset);
@@ -239,7 +230,7 @@ struct GroupSlicer {
                        return idx - static_cast<int64_t>(offset);
                      });
 
-      std::decay_t<A1> typedTable{{groupedElementsTable}, std::move(slicedSelection), offset};
+      std::decay_t<A1> typedTable{{groupedElementsTable}, std::move(slicedSelection)};
       typedTable.bindInternalIndicesTo(&originalTable);
       return typedTable;
     }
@@ -281,9 +272,6 @@ struct GroupSlicer {
     std::span<int64_t const> groupSelection;
     std::array<std::span<int64_t const> const*, sizeof...(A)> selections;
     std::array<std::span<int64_t const>::iterator, sizeof...(A)> starts;
-    // Cached empty (0-row) table per associated table, lazily built and reused
-    // for empty groups so we do not slice every column on each empty group.
-    std::array<std::shared_ptr<arrow::Table>, sizeof...(A)> emptyTables{};
 
     std::array<SliceInfoPtr, sizeof...(A)> sliceInfos;
     std::array<SliceInfoUnsortedPtr, sizeof...(A)> sliceInfosUnsorted;

Framework/Core/include/Framework/GroupedCombinations.h

@@ -70,15 +70,15 @@ struct GroupedCombinationsGenerator {
     template <typename... T2s>
     GroupedIterator(const GroupingPolicy& groupingPolicy, const G& grouping, const std::tuple<T2s...>& associated, SliceCache* cache_)
       : GroupingPolicy(groupingPolicy),
-        mGrouping{std::make_shared<G>(std::vector{grouping.asArrowTable()})},
+        mGrouping{std::make_shared<G>(std::vector{grouping.asArrowTableRef()})},
         mAssociated{std::make_shared<std::tuple<As...>>(std::make_tuple(std::get<has_type_at<As>(pack<T2s...>{})>(associated)...))},
         mIndexColumns{getMatchingIndexNode<G, As>()...},
         cache{cache_}
     {
       if constexpr (soa::is_filtered_table<std::decay_t<G>>) {
-        mGrouping = std::make_shared<G>(std::vector{grouping.asArrowTable()}, grouping.getSelectedRows());
+        mGrouping = std::make_shared<G>(std::vector{grouping.asArrowTableRef()}, grouping.getSelectedRows());
       } else {
-        mGrouping = std::make_shared<G>(std::vector{grouping.asArrowTable()});
+        mGrouping = std::make_shared<G>(std::vector{grouping.asArrowTableRef()});
       }
       setMultipleGroupingTables<sizeof...(As)>(grouping);
       if (!this->mIsEnd) {
@@ -94,9 +94,9 @@ struct GroupedCombinationsGenerator {
     void setTables(const G& grouping, const std::tuple<T2s...>& associated)
     {
       if constexpr (soa::is_filtered_table<std::decay_t<G>>) {
-        mGrouping = std::make_shared<G>(std::vector{grouping.asArrowTable()}, grouping.getSelectedRows());
+        mGrouping = std::make_shared<G>(std::vector{grouping.asArrowTableRef()}, grouping.getSelectedRows());
       } else {
-        mGrouping = std::make_shared<G>(std::vector{grouping.asArrowTable()});
+        mGrouping = std::make_shared<G>(std::vector{grouping.asArrowTableRef()});
       }
       mAssociated = std::make_shared<std::tuple<As...>>(std::make_tuple(std::get<has_type_at_v<As>(pack<T2s...>{})>(associated)...));
       setMultipleGroupingTables<sizeof...(As)>(grouping);