Skip to content

Commit 94e1eec

Browse files
aalkinaalkin
authored
DPL Analysis: avoid Arrow's Slice API (#15613)
1 parent e6152d7 commit 94e1eec

16 files changed

Lines changed: 446 additions & 413 deletions

‎Framework/Core/include/Framework/ASoA.h‎

Lines changed: 204 additions & 276 deletions
Large diffs are not rendered by default.

‎Framework/Core/include/Framework/AnalysisHelpers.h‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -147,7 +147,7 @@ auto spawner(framework::pack<C...>, std::vector<std::shared_ptr<arrow::Table>>&&
147147
if (fullTable->num_rows() == 0) {
148148
return makeEmptyTable(name, framework::pack<C...>{});
149149
}
150-
return spawnerHelper(fullTable, schema, sizeof...(C), projectors, name, projector);
150+
return spawnerHelper(fullTable.tablePtr, schema, sizeof...(C), projectors, name, projector);
151151
}
152152

153153
std::string serializeProjectors(std::vector<framework::expressions::Projector>& projectors);
@@ -950,7 +950,7 @@ auto getTableFromFilter(soa::is_filtered_table auto const& table, soa::Selection
950950

951951
auto getTableFromFilter(soa::is_not_filtered_table auto const& table, soa::SelectionVector&& selection)
952952
{
953-
return std::make_unique<o2::soa::Filtered<std::decay_t<decltype(table)>>>(std::vector{table.asArrowTable()}, std::forward<soa::SelectionVector>(selection));
953+
return std::make_unique<o2::soa::Filtered<std::decay_t<decltype(table)>>>(std::vector{table.asArrowTableRef()}, std::forward<soa::SelectionVector>(selection));
954954
}
955955

956956
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 {
982982

983983
void bindTable(T const& table)
984984
{
985-
intializeCaches(T::table_t::hashes(), table.asArrowTable()->schema());
985+
intializeCaches(T::table_t::hashes(), table.asArrowTableRef()->schema());
986986
if (dataframeChanged) {
987987
mFiltered = getTableFromFilter(table, soa::selectionToVector(framework::expressions::createSelection(table.asArrowTable(), gfilter)));
988988
dataframeChanged = false;
@@ -1086,7 +1086,7 @@ auto Extend(T const& table)
10861086
static std::array<framework::expressions::Projector, sizeof...(Cs)> projectors{{std::move(Cs::Projector())...}};
10871087
static std::shared_ptr<gandiva::Projector> projector = nullptr;
10881088
static auto schema = std::make_shared<arrow::Schema>(o2::soa::createFieldsFromColumns(framework::pack<Cs...>{}));
1089-
return output_t{{o2::framework::spawner(framework::pack<Cs...>{}, {table.asArrowTable()}, "dynamicExtension", projectors.data(), projector, schema), table.asArrowTable()}, 0};
1089+
return output_t{{o2::framework::spawner(framework::pack<Cs...>{}, {table.asArrowTable()}, "dynamicExtension", projectors.data(), projector, schema), table.asArrowTable()}};
10901090
}
10911091

10921092
/// 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>
10951095
auto Attach(T const& table)
10961096
{
10971097
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...>>;
1098-
return output_t{{table.asArrowTable()}, table.offset()};
1098+
return output_t{{table.asArrowTableRef()}};
10991099
}
11001100
} // namespace o2::soa
11011101

‎Framework/Core/include/Framework/AnalysisManagers.h‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -319,12 +319,12 @@ bool prepareOutput(ProcessingContext& context, T& spawns)
319319
}
320320
using D = o2::aod::Hash<metadata::extension_table_t::ref.desc_hash>;
321321

322-
spawns.extension = std::make_shared<typename T::extension_t>(o2::framework::spawner<D>(originalTable,
322+
spawns.extension = std::make_shared<typename T::extension_t>(o2::framework::spawner<D>(originalTable.tablePtr,
323323
o2::aod::label<metadata::extension_table_t::ref>(),
324324
spawns.projectors.data(),
325325
spawns.projector,
326326
spawns.schema));
327-
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}));
327+
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}));
328328
return true;
329329
}
330330

@@ -348,12 +348,12 @@ bool prepareOutput(ProcessingContext& context, T& defines)
348348
}
349349
using D = o2::aod::Hash<metadata::extension_table_t::ref.desc_hash>;
350350

351-
defines.extension = std::make_shared<typename T::extension_t>(o2::framework::spawner<D>(originalTable,
351+
defines.extension = std::make_shared<typename T::extension_t>(o2::framework::spawner<D>(originalTable.tablePtr,
352352
o2::aod::label<metadata::extension_table_t::ref>(),
353353
defines.projectors.data(),
354354
defines.projector,
355355
defines.schema));
356-
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}));
356+
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}));
357357
return true;
358358
}
359359

@@ -380,12 +380,12 @@ bool prepareDelayedOutput(ProcessingContext& context, T& defines)
380380
}
381381
using D = o2::aod::Hash<metadata::extension_table_t::ref.desc_hash>;
382382

383-
defines.extension = std::make_shared<typename T::extension_t>(o2::framework::spawner<D>(originalTable,
383+
defines.extension = std::make_shared<typename T::extension_t>(o2::framework::spawner<D>(originalTable.tablePtr,
384384
o2::aod::label<metadata::extension_table_t::ref>(),
385385
defines.projectors.data(),
386386
defines.projector,
387387
defines.schema));
388-
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}));
388+
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}));
389389
return true;
390390
}
391391

‎Framework/Core/include/Framework/AnalysisTask.h‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -226,7 +226,7 @@ struct AnalysisDataProcessorBuilder {
226226
template <std::ranges::input_range R>
227227
static auto extractTablesFromRecord(InputRecord& record, R matchers)
228228
{
229-
std::vector<std::shared_ptr<arrow::Table>> tables;
229+
std::vector<soa::ArrowTableRef> tables;
230230
std::ranges::transform(matchers, std::back_inserter(tables), [&record](auto const& m) {
231231
return record.get<TableConsumer>(m.second)->asArrowTable();
232232
});
@@ -248,8 +248,8 @@ struct AnalysisDataProcessorBuilder {
248248
template <soa::is_filtered T, std::ranges::input_range R>
249249
static auto extractFilteredFromRecord(InputRecord& record, R matchers, ExpressionInfo& info)
250250
{
251-
std::shared_ptr<arrow::Table> table = soa::ArrowHelpers::joinTables(extractTablesFromRecord(record, matchers));
252-
expressions::updateFilterInfo(info, table);
251+
auto table = soa::ArrowHelpers::joinTables(extractTablesFromRecord(record, matchers));
252+
expressions::updateFilterInfo(info, table.tablePtr);
253253
if constexpr (!o2::soa::is_smallgroups<std::decay_t<T>>) {
254254
if (info.selection == nullptr) {
255255
soa::missingFilterDeclaration(info.processHash, info.argumentIndex);

‎Framework/Core/include/Framework/ArrowTypes.h‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,12 +11,55 @@
1111

1212
#ifndef O2_FRAMEWORK_ARROWTYPES_H
1313
#define O2_FRAMEWORK_ARROWTYPES_H
14+
#include <arrow/table.h>
1415
#include "Framework/Traits.h"
1516
#include "arrow/type_fwd.h"
1617
#include <span>
1718

1819
namespace o2::soa
1920
{
21+
struct ArrowRange {
22+
uint64_t offset;
23+
int64_t size;
24+
25+
bool operator!=(ArrowRange const& other) const
26+
{
27+
return (offset != other.offset) && (size != other.size);
28+
}
29+
};
30+
31+
struct ArrowTableRef {
32+
std::shared_ptr<arrow::Table> tablePtr = nullptr;
33+
ArrowRange range{0, 0};
34+
35+
ArrowTableRef() = default;
36+
ArrowTableRef(std::shared_ptr<arrow::Table> table)
37+
: tablePtr{table},
38+
range{0, table->num_rows()}
39+
{
40+
}
41+
ArrowTableRef(std::shared_ptr<arrow::Table> table, ArrowRange range_)
42+
: tablePtr{table},
43+
range{range_}
44+
{
45+
}
46+
47+
ArrowTableRef makeEmpty() const
48+
{
49+
return {tablePtr, {0, 0}};
50+
}
51+
52+
ArrowTableRef slice(ArrowRange newRange) const
53+
{
54+
return {tablePtr, newRange};
55+
}
56+
57+
std::shared_ptr<arrow::Table> const& operator->() const
58+
{
59+
return tablePtr;
60+
}
61+
};
62+
2063
template <typename T>
2164
struct arrow_array_for {
2265
};

‎Framework/Core/include/Framework/GroupSlicer.h‎

Lines changed: 3 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -194,7 +194,7 @@ struct GroupSlicer {
194194
}
195195
}
196196
}
197-
std::decay_t<A1> typedTable{{originalTable.asArrowTable()}, std::move(s)};
197+
std::decay_t<A1> typedTable{{originalTable.asArrowTableRef()}, std::move(s)};
198198
typedTable.bindInternalIndicesTo(&originalTable);
199199
return typedTable;
200200
}
@@ -218,16 +218,7 @@ struct GroupSlicer {
218218
auto oc = sliceInfos[index].getSliceFor(pos);
219219
uint64_t offset = oc.first;
220220
auto count = oc.second;
221-
if (count == 0) {
222-
// Empty group: avoid slicing every column only to discard it. Cache one
223-
// empty (0-row) table per associated table and reuse it. This is the
224-
// common case for sparse grouping (e.g. collisions with no candidates).
225-
if (!emptyTables[index]) {
226-
emptyTables[index] = originalTable.asArrowTable()->Slice(0, 0);
227-
}
228-
return std::decay_t<A1>{{emptyTables[index]}, soa::SelectionVector{}};
229-
}
230-
auto groupedElementsTable = originalTable.asArrowTable()->Slice(offset, count);
221+
auto groupedElementsTable = originalTable.asArrowTableRef().slice({offset, count});
231222

232223
// for each grouping element we need to slice the selection vector
233224
auto start_iterator = std::lower_bound(starts[index], selections[index]->end(), offset);
@@ -239,7 +230,7 @@ struct GroupSlicer {
239230
return idx - static_cast<int64_t>(offset);
240231
});
241232

242-
std::decay_t<A1> typedTable{{groupedElementsTable}, std::move(slicedSelection), offset};
233+
std::decay_t<A1> typedTable{{groupedElementsTable}, std::move(slicedSelection)};
243234
typedTable.bindInternalIndicesTo(&originalTable);
244235
return typedTable;
245236
}
@@ -281,9 +272,6 @@ struct GroupSlicer {
281272
std::span<int64_t const> groupSelection;
282273
std::array<std::span<int64_t const> const*, sizeof...(A)> selections;
283274
std::array<std::span<int64_t const>::iterator, sizeof...(A)> starts;
284-
// Cached empty (0-row) table per associated table, lazily built and reused
285-
// for empty groups so we do not slice every column on each empty group.
286-
std::array<std::shared_ptr<arrow::Table>, sizeof...(A)> emptyTables{};
287275

288276
std::array<SliceInfoPtr, sizeof...(A)> sliceInfos;
289277
std::array<SliceInfoUnsortedPtr, sizeof...(A)> sliceInfosUnsorted;

‎Framework/Core/include/Framework/GroupedCombinations.h‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -70,15 +70,15 @@ struct GroupedCombinationsGenerator {
7070
template <typename... T2s>
7171
GroupedIterator(const GroupingPolicy& groupingPolicy, const G& grouping, const std::tuple<T2s...>& associated, SliceCache* cache_)
7272
: GroupingPolicy(groupingPolicy),
73-
mGrouping{std::make_shared<G>(std::vector{grouping.asArrowTable()})},
73+
mGrouping{std::make_shared<G>(std::vector{grouping.asArrowTableRef()})},
7474
mAssociated{std::make_shared<std::tuple<As...>>(std::make_tuple(std::get<has_type_at<As>(pack<T2s...>{})>(associated)...))},
7575
mIndexColumns{getMatchingIndexNode<G, As>()...},
7676
cache{cache_}
7777
{
7878
if constexpr (soa::is_filtered_table<std::decay_t<G>>) {
79-
mGrouping = std::make_shared<G>(std::vector{grouping.asArrowTable()}, grouping.getSelectedRows());
79+
mGrouping = std::make_shared<G>(std::vector{grouping.asArrowTableRef()}, grouping.getSelectedRows());
8080
} else {
81-
mGrouping = std::make_shared<G>(std::vector{grouping.asArrowTable()});
81+
mGrouping = std::make_shared<G>(std::vector{grouping.asArrowTableRef()});
8282
}
8383
setMultipleGroupingTables<sizeof...(As)>(grouping);
8484
if (!this->mIsEnd) {
@@ -94,9 +94,9 @@ struct GroupedCombinationsGenerator {
9494
void setTables(const G& grouping, const std::tuple<T2s...>& associated)
9595
{
9696
if constexpr (soa::is_filtered_table<std::decay_t<G>>) {
97-
mGrouping = std::make_shared<G>(std::vector{grouping.asArrowTable()}, grouping.getSelectedRows());
97+
mGrouping = std::make_shared<G>(std::vector{grouping.asArrowTableRef()}, grouping.getSelectedRows());
9898
} else {
99-
mGrouping = std::make_shared<G>(std::vector{grouping.asArrowTable()});
99+
mGrouping = std::make_shared<G>(std::vector{grouping.asArrowTableRef()});
100100
}
101101
mAssociated = std::make_shared<std::tuple<As...>>(std::make_tuple(std::get<has_type_at_v<As>(pack<T2s...>{})>(associated)...));
102102
setMultipleGroupingTables<sizeof...(As)>(grouping);

0 commit comments

Comments
 (0)