Skip to content

Commit 0f1ed00

Browse files
aalkinaalkin
authored
DPL Analysis: Refactor group slicer (#13839)
1 parent 5efb6e2 commit 0f1ed00

File tree

1 file changed

+116
-85
lines changed

1 file changed

+116
-85
lines changed

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

Lines changed: 116 additions & 85 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,6 @@
2222

2323
namespace o2::framework
2424
{
25-
2625
template <typename G, typename... A>
2726
struct GroupSlicer {
2827
using grouping_t = std::decay_t<G>;
@@ -46,34 +45,41 @@ struct GroupSlicer {
4645
GroupSlicerIterator& operator=(GroupSlicerIterator&&) = default;
4746

4847
template <typename T>
48+
auto splittingFunction(T&&)
49+
{
50+
}
51+
52+
template <typename T>
53+
requires(o2::soa::relatedByIndex<std::decay_t<G>, std::decay_t<T>>())
4954
auto splittingFunction(T&& table)
5055
{
5156
constexpr auto index = framework::has_type_at_v<std::decay_t<T>>(associated_pack_t{});
52-
if constexpr (o2::soa::relatedByIndex<std::decay_t<G>, std::decay_t<T>>()) {
53-
auto binding = o2::soa::getLabelFromTypeForKey<std::decay_t<T>>(mIndexColumnName);
54-
auto bk = std::make_pair(binding, mIndexColumnName);
55-
if constexpr (!o2::soa::is_smallgroups<std::decay_t<T>>) {
56-
if (table.size() == 0) {
57-
return;
58-
}
59-
sliceInfos[index] = mSlices->getCacheFor(bk);
60-
} else {
61-
if (table.tableSize() == 0) {
62-
return;
63-
}
64-
sliceInfosUnsorted[index] = mSlices->getCacheUnsortedFor(bk);
57+
auto binding = o2::soa::getLabelFromTypeForKey<std::decay_t<T>>(mIndexColumnName);
58+
auto bk = std::make_pair(binding, mIndexColumnName);
59+
if constexpr (!o2::soa::is_smallgroups<std::decay_t<T>>) {
60+
if (table.size() == 0) {
61+
return;
6562
}
63+
sliceInfos[index] = mSlices->getCacheFor(bk);
64+
} else {
65+
if (table.tableSize() == 0) {
66+
return;
67+
}
68+
sliceInfosUnsorted[index] = mSlices->getCacheUnsortedFor(bk);
6669
}
6770
}
6871

6972
template <typename T>
73+
auto extractingFunction(T&&)
74+
{
75+
}
76+
77+
template <soa::is_filtered_table T>
7078
auto extractingFunction(T&& table)
7179
{
72-
if constexpr (soa::is_filtered_table<std::decay_t<T>>) {
73-
constexpr auto index = framework::has_type_at_v<std::decay_t<T>>(associated_pack_t{});
74-
selections[index] = &table.getSelectedRows();
75-
starts[index] = selections[index]->begin();
76-
}
80+
constexpr auto index = framework::has_type_at_v<std::decay_t<T>>(associated_pack_t{});
81+
selections[index] = &table.getSelectedRows();
82+
starts[index] = selections[index]->begin();
7783
}
7884

7985
GroupSlicerIterator(G& gt, std::tuple<A...>& at, ArrowTableSlicingCache& slices)
@@ -151,83 +157,108 @@ struct GroupSlicer {
151157
return std::make_tuple(prepareArgument<A>()...);
152158
}
153159

154-
template <typename A1>
160+
template <soa::is_smallgroups A1>
161+
requires(o2::soa::relatedByIndex<std::decay_t<G>, std::decay_t<A1>>() && soa::is_filtered_table<A1>)
155162
auto prepareArgument()
156163
{
157164
constexpr auto index = framework::has_type_at_v<A1>(associated_pack_t{});
158165
auto& originalTable = std::get<A1>(*mAt);
159-
160-
if constexpr (o2::soa::relatedByIndex<std::decay_t<G>, std::decay_t<A1>>()) {
161-
uint64_t pos;
162-
if constexpr (soa::is_filtered_table<std::decay_t<G>>) {
163-
pos = groupSelection[position];
164-
} else {
165-
pos = position;
166+
uint64_t pos;
167+
if constexpr (soa::is_filtered_table<std::decay_t<G>>) {
168+
pos = groupSelection[position];
169+
} else {
170+
pos = position;
171+
}
172+
// generic split
173+
auto selection = sliceInfosUnsorted[index].getSliceFor(pos);
174+
// intersect selections
175+
o2::soa::SelectionVector s;
176+
if (selections[index]->empty()) {
177+
if (!selection.empty()) {
178+
std::copy(selection.begin(), selection.end(), std::back_inserter(s));
166179
}
167-
168-
if constexpr (!o2::soa::is_smallgroups<std::decay_t<A1>>) {
169-
// optimized split
170-
if (originalTable.size() == 0) {
171-
return originalTable;
172-
}
173-
auto oc = sliceInfos[index].getSliceFor(pos);
174-
uint64_t offset = oc.first;
175-
auto count = oc.second;
176-
if constexpr (soa::is_filtered_table<std::decay_t<A1>>) {
177-
auto groupedElementsTable = originalTable.asArrowTable()->Slice(offset, count);
178-
if (count == 0) {
179-
return std::decay_t<A1>{{groupedElementsTable}, soa::SelectionVector{}};
180-
}
181-
182-
// for each grouping element we need to slice the selection vector
183-
auto start_iterator = std::lower_bound(starts[index], selections[index]->end(), offset);
184-
auto stop_iterator = std::lower_bound(start_iterator, selections[index]->end(), offset + count);
185-
starts[index] = stop_iterator;
186-
soa::SelectionVector slicedSelection{start_iterator, stop_iterator};
187-
std::transform(slicedSelection.begin(), slicedSelection.end(), slicedSelection.begin(),
188-
[&offset](int64_t idx) {
189-
return idx - static_cast<int64_t>(offset);
190-
});
191-
192-
std::decay_t<A1> typedTable{{groupedElementsTable}, std::move(slicedSelection), offset};
193-
typedTable.bindInternalIndicesTo(&originalTable);
194-
return typedTable;
195-
196-
} else {
197-
auto groupedElementsTable = originalTable.rawSlice(offset, offset + count - 1);
198-
groupedElementsTable.bindInternalIndicesTo(&originalTable);
199-
return groupedElementsTable;
200-
}
201-
} else {
202-
// generic split
203-
if constexpr (soa::is_filtered_table<std::decay_t<A1>>) {
204-
auto selection = sliceInfosUnsorted[index].getSliceFor(pos);
205-
// intersect selections
206-
o2::soa::SelectionVector s;
207-
if (selections[index]->empty()) {
208-
if (!selection.empty()) {
209-
std::copy(selection.begin(), selection.end(), std::back_inserter(s));
210-
}
211-
} else {
212-
if (!selection.empty()) {
213-
if constexpr (std::decay_t<A1>::applyFilters) {
214-
std::set_intersection(selection.begin(), selection.end(), selections[index]->begin(), selections[index]->end(), std::back_inserter(s));
215-
} else {
216-
std::copy(selection.begin(), selection.end(), std::back_inserter(s));
217-
}
218-
}
219-
}
220-
std::decay_t<A1> typedTable{{originalTable.asArrowTable()}, std::move(s)};
221-
typedTable.bindInternalIndicesTo(&originalTable);
222-
return typedTable;
180+
} else {
181+
if (!selection.empty()) {
182+
if constexpr (std::decay_t<A1>::applyFilters) {
183+
std::set_intersection(selection.begin(), selection.end(), selections[index]->begin(), selections[index]->end(), std::back_inserter(s));
223184
} else {
224-
throw runtime_error("Unsorted grouped table needs to be used with soa::SmallGroups<>");
185+
std::copy(selection.begin(), selection.end(), std::back_inserter(s));
225186
}
226187
}
188+
}
189+
std::decay_t<A1> typedTable{{originalTable.asArrowTable()}, std::move(s)};
190+
typedTable.bindInternalIndicesTo(&originalTable);
191+
return typedTable;
192+
}
193+
194+
template <soa::is_filtered_table A1>
195+
requires(o2::soa::relatedByIndex<std::decay_t<G>, std::decay_t<A1>>() && !soa::is_smallgroups<A1>)
196+
auto prepareArgument()
197+
{
198+
constexpr auto index = framework::has_type_at_v<A1>(associated_pack_t{});
199+
auto& originalTable = std::get<A1>(*mAt);
200+
if (originalTable.size() == 0) {
201+
return originalTable;
202+
}
203+
uint64_t pos;
204+
if constexpr (soa::is_filtered_table<std::decay_t<G>>) {
205+
pos = groupSelection[position];
227206
} else {
228-
static_assert(!o2::soa::is_smallgroups<std::decay_t<A1>>, "SmallGroups used with a table that is not related by index to the gouping table");
207+
pos = position;
208+
}
209+
// optimized split
210+
auto oc = sliceInfos[index].getSliceFor(pos);
211+
uint64_t offset = oc.first;
212+
auto count = oc.second;
213+
auto groupedElementsTable = originalTable.asArrowTable()->Slice(offset, count);
214+
if (count == 0) {
215+
return std::decay_t<A1>{{groupedElementsTable}, soa::SelectionVector{}};
216+
}
217+
218+
// for each grouping element we need to slice the selection vector
219+
auto start_iterator = std::lower_bound(starts[index], selections[index]->end(), offset);
220+
auto stop_iterator = std::lower_bound(start_iterator, selections[index]->end(), offset + count);
221+
starts[index] = stop_iterator;
222+
soa::SelectionVector slicedSelection{start_iterator, stop_iterator};
223+
std::transform(slicedSelection.begin(), slicedSelection.end(), slicedSelection.begin(),
224+
[&offset](int64_t idx) {
225+
return idx - static_cast<int64_t>(offset);
226+
});
227+
228+
std::decay_t<A1> typedTable{{groupedElementsTable}, std::move(slicedSelection), offset};
229+
typedTable.bindInternalIndicesTo(&originalTable);
230+
return typedTable;
231+
}
232+
233+
template <soa::is_table A1>
234+
requires(o2::soa::relatedByIndex<std::decay_t<G>, std::decay_t<A1>>() && !soa::is_smallgroups<A1> && !soa::is_filtered_table<A1>)
235+
auto prepareArgument()
236+
{
237+
constexpr auto index = framework::has_type_at_v<A1>(associated_pack_t{});
238+
auto& originalTable = std::get<A1>(*mAt);
239+
if (originalTable.size() == 0) {
229240
return originalTable;
230241
}
242+
uint64_t pos;
243+
if constexpr (soa::is_filtered_table<std::decay_t<G>>) {
244+
pos = groupSelection[position];
245+
} else {
246+
pos = position;
247+
}
248+
// optimized split
249+
auto oc = sliceInfos[index].getSliceFor(pos);
250+
uint64_t offset = oc.first;
251+
auto count = oc.second;
252+
auto groupedElementsTable = originalTable.rawSlice(offset, offset + count - 1);
253+
groupedElementsTable.bindInternalIndicesTo(&originalTable);
254+
return groupedElementsTable;
255+
}
256+
257+
template <soa::is_table A1>
258+
requires(!o2::soa::relatedByIndex<std::decay_t<G>, std::decay_t<A1>>() && !soa::is_smallgroups<A1>)
259+
auto prepareArgument()
260+
{
261+
return std::get<A1>(*mAt);
231262
}
232263

233264
std::string mIndexColumnName;

0 commit comments

Comments
 (0)