Skip to content

Commit ee5b059

Browse files
aalkinaalkin
committed
DPL Analysis: refactor Preslice; move some code out of line
1 parent d26f229 commit ee5b059

File tree

4 files changed

+191
-136
lines changed

4 files changed

+191
-136
lines changed

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

Lines changed: 119 additions & 113 deletions
Original file line numberDiff line numberDiff line change
@@ -1389,76 +1389,69 @@ consteval static bool relatedBySortedIndex()
13891389

13901390
namespace o2::framework
13911391
{
1392-
template <typename T, bool OPT = false, bool SORTED = true>
1393-
struct PresliceBase {
1394-
constexpr static bool sorted = SORTED;
1392+
1393+
struct PreslicePolicyBase {
1394+
const std::string binding;
1395+
StringPair bindingKey;
1396+
1397+
bool isMissing() const;
1398+
StringPair const& getBindingKey() const;
1399+
};
1400+
1401+
struct PreslicePolicySorted : public PreslicePolicyBase {
1402+
void updateSliceInfo(SliceInfoPtr&& si);
1403+
1404+
SliceInfoPtr sliceInfo;
1405+
std::shared_ptr<arrow::Table> getSliceFor(int value, std::shared_ptr<arrow::Table> const& input, uint64_t& offset) const;
1406+
};
1407+
1408+
struct PreslicePolicyGeneral : public PreslicePolicyBase {
1409+
void updateSliceInfo(SliceInfoUnsortedPtr&& si);
1410+
1411+
SliceInfoUnsortedPtr sliceInfo;
1412+
gsl::span<const int64_t> getSliceFor(int value) const;
1413+
};
1414+
1415+
template <typename T, typename Policy, bool OPT = false>
1416+
struct PresliceBase : public Policy {
1417+
constexpr static bool sorted = std::same_as<Policy, PreslicePolicySorted>;
13951418
constexpr static bool optional = OPT;
13961419
using target_t = T;
13971420
const std::string binding;
13981421

13991422
PresliceBase(expressions::BindingNode index_)
1400-
: binding{o2::soa::getLabelFromTypeForKey<T, OPT>(index_.name)},
1401-
bindingKey{binding, index_.name} {}
1402-
1403-
void updateSliceInfo(std::conditional_t<SORTED, SliceInfoPtr, SliceInfoUnsortedPtr>&& si)
1404-
{
1405-
sliceInfo = si;
1406-
}
1423+
: Policy{PreslicePolicyBase{{o2::soa::getLabelFromTypeForKey<T, OPT>(index_.name)},{binding, index_.name}}, {}}
1424+
{}
14071425

14081426
std::shared_ptr<arrow::Table> getSliceFor(int value, std::shared_ptr<arrow::Table> const& input, uint64_t& offset) const
14091427
{
14101428
if constexpr (OPT) {
1411-
if (isMissing()) {
1429+
if (Policy::isMissing()) {
14121430
return nullptr;
14131431
}
14141432
}
1415-
if constexpr (SORTED) {
1416-
auto [offset_, count] = sliceInfo.getSliceFor(value);
1417-
auto output = input->Slice(offset_, count);
1418-
offset = static_cast<int64_t>(offset_);
1419-
return output;
1420-
} else {
1421-
static_assert(SORTED, "Wrong method called for unsorted cache");
1422-
}
1433+
return Policy::getSliceFor(value, input, offset);
14231434
}
14241435

14251436
gsl::span<const int64_t> getSliceFor(int value) const
14261437
{
14271438
if constexpr (OPT) {
1428-
if (isMissing()) {
1439+
if (Policy::isMissing()) {
14291440
return {};
14301441
}
14311442
}
1432-
if constexpr (!SORTED) {
1433-
return sliceInfo.getSliceFor(value);
1434-
} else {
1435-
static_assert(!SORTED, "Wrong method called for sorted cache");
1436-
}
1443+
return Policy::getSliceFor(value);
14371444
}
1438-
1439-
bool isMissing() const
1440-
{
1441-
return binding == "[MISSING]";
1442-
}
1443-
1444-
StringPair const& getBindingKey() const
1445-
{
1446-
return bindingKey;
1447-
}
1448-
1449-
std::conditional_t<SORTED, SliceInfoPtr, SliceInfoUnsortedPtr> sliceInfo;
1450-
1451-
StringPair bindingKey;
14521445
};
14531446

14541447
template <typename T>
1455-
using PresliceUnsorted = PresliceBase<T, false, false>;
1448+
using PresliceUnsorted = PresliceBase<T, PreslicePolicyGeneral, false>;
14561449
template <typename T>
1457-
using PresliceUnsortedOptional = PresliceBase<T, true, false>;
1450+
using PresliceUnsortedOptional = PresliceBase<T, PreslicePolicyGeneral, true>;
14581451
template <typename T>
1459-
using Preslice = PresliceBase<T, false, true>;
1452+
using Preslice = PresliceBase<T, PreslicePolicySorted, false>;
14601453
template <typename T>
1461-
using PresliceOptional = PresliceBase<T, true, true>;
1454+
using PresliceOptional = PresliceBase<T, PreslicePolicySorted, true>;
14621455

14631456
} // namespace o2::framework
14641457

@@ -1497,44 +1490,54 @@ static consteval auto extractBindings(framework::pack<Is...>)
14971490

14981491
SelectionVector selectionToVector(gandiva::Selection const& sel);
14991492

1500-
template <typename T, typename C, bool OPT, bool SORTED>
1501-
auto doSliceBy(T const* table, o2::framework::PresliceBase<C, OPT, SORTED> const& container, int value)
1493+
template <typename T, typename C, typename Policy, bool OPT>
1494+
requires std::same_as<Policy, framework::PreslicePolicySorted> && (o2::soa::is_binding_compatible_v<C, T>())
1495+
auto doSliceBy(T const* table, o2::framework::PresliceBase<C, Policy, OPT> const& container, int value)
15021496
{
1503-
if constexpr (o2::soa::is_binding_compatible_v<C, T>()) {
1504-
if constexpr (OPT) {
1505-
if (container.isMissing()) {
1506-
missingOptionalPreslice(getLabelFromType<std::decay_t<T>>().data(), container.bindingKey.second.c_str());
1507-
}
1508-
}
1509-
if constexpr (SORTED) {
1510-
uint64_t offset = 0;
1511-
auto out = container.getSliceFor(value, table->asArrowTable(), offset);
1512-
auto t = typename T::self_t({out}, offset);
1513-
table->copyIndexBindings(t);
1514-
t.bindInternalIndicesTo(table);
1515-
return t;
1516-
} else {
1517-
auto selection = container.getSliceFor(value);
1518-
if constexpr (soa::is_filtered_table<T>) {
1519-
auto t = soa::Filtered<typename T::base_t>({table->asArrowTable()}, selection);
1520-
table->copyIndexBindings(t);
1521-
t.bindInternalIndicesTo(table);
1522-
t.intersectWithSelection(table->getSelectedRows()); // intersect filters
1523-
return t;
1524-
} else {
1525-
auto t = soa::Filtered<T>({table->asArrowTable()}, selection);
1526-
table->copyIndexBindings(t);
1527-
t.bindInternalIndicesTo(table);
1528-
return t;
1529-
}
1497+
if constexpr (OPT) {
1498+
if (container.isMissing()) {
1499+
missingOptionalPreslice(getLabelFromType<std::decay_t<T>>().data(), container.bindingKey.second.c_str());
15301500
}
1531-
} else {
1532-
if constexpr (SORTED) {
1533-
static_assert(o2::framework::always_static_assert_v<C>, "Wrong Preslice<> entry used: incompatible type");
1534-
} else {
1535-
static_assert(o2::framework::always_static_assert_v<C>, "Wrong PresliceUnsorted<> entry used: incompatible type");
1501+
}
1502+
uint64_t offset = 0;
1503+
auto out = container.getSliceFor(value, table->asArrowTable(), offset);
1504+
auto t = typename T::self_t({out}, offset);
1505+
table->copyIndexBindings(t);
1506+
t.bindInternalIndicesTo(table);
1507+
return t;
1508+
}
1509+
1510+
template <soa::is_filtered_table T>
1511+
auto doSliceByHelper(T const* table, gsl::span<const int64_t> const& selection)
1512+
{
1513+
auto t = soa::Filtered<typename T::base_t>({table->asArrowTable()}, selection);
1514+
table->copyIndexBindings(t);
1515+
t.bindInternalIndicesTo(table);
1516+
t.intersectWithSelection(table->getSelectedRows()); // intersect filters
1517+
return t;
1518+
}
1519+
1520+
template <soa::is_table T>
1521+
requires (!soa::is_filtered_table<T>)
1522+
auto doSliceByHelper(T const* table, gsl::span<const int64_t> const& selection)
1523+
{
1524+
auto t = soa::Filtered<T>({table->asArrowTable()}, selection);
1525+
table->copyIndexBindings(t);
1526+
t.bindInternalIndicesTo(table);
1527+
return t;
1528+
}
1529+
1530+
template <typename T, typename C, typename Policy, bool OPT>
1531+
requires std::same_as<Policy, framework::PreslicePolicyGeneral> && (o2::soa::is_binding_compatible_v<C, T>())
1532+
auto doSliceBy(T const* table, o2::framework::PresliceBase<C, Policy, OPT> const& container, int value)
1533+
{
1534+
if constexpr (OPT) {
1535+
if (container.isMissing()) {
1536+
missingOptionalPreslice(getLabelFromType<std::decay_t<T>>().data(), container.bindingKey.second.c_str());
15361537
}
15371538
}
1539+
auto selection = container.getSliceFor(value);
1540+
return doSliceByHelper(table, selection);
15381541
}
15391542

15401543
template <typename T>
@@ -1573,20 +1576,17 @@ auto prepareFilteredSlice(T const* table, std::shared_ptr<arrow::Table> slice, u
15731576
}
15741577

15751578
template <typename T, typename C, bool OPT>
1576-
auto doFilteredSliceBy(T const* table, o2::framework::PresliceBase<C, OPT> const& container, int value)
1579+
requires (o2::soa::is_binding_compatible_v<C, T>())
1580+
auto doFilteredSliceBy(T const* table, o2::framework::PresliceBase<C, framework::PreslicePolicySorted, OPT> const& container, int value)
15771581
{
1578-
if constexpr (o2::soa::is_binding_compatible_v<C, T>()) {
1579-
if constexpr (OPT) {
1580-
if (container.isMissing()) {
1581-
missingOptionalPreslice(getLabelFromType<T>().data(), container.bindingKey.second.c_str());
1582-
}
1582+
if constexpr (OPT) {
1583+
if (container.isMissing()) {
1584+
missingOptionalPreslice(getLabelFromType<T>().data(), container.bindingKey.second.c_str());
15831585
}
1584-
uint64_t offset = 0;
1585-
auto slice = container.getSliceFor(value, table->asArrowTable(), offset);
1586-
return prepareFilteredSlice(table, slice, offset);
1587-
} else {
1588-
static_assert(o2::framework::always_static_assert_v<C>, "Wrong Preslice<> entry used: incompatible type");
15891586
}
1587+
uint64_t offset = 0;
1588+
auto slice = container.getSliceFor(value, table->asArrowTable(), offset);
1589+
return prepareFilteredSlice(table, slice, offset);
15901590
}
15911591

15921592
template <typename T>
@@ -2099,8 +2099,8 @@ class Table
20992099
return doSliceByCachedUnsorted(this, node, value, cache);
21002100
}
21012101

2102-
template <typename T1, bool OPT, bool SORTED>
2103-
auto sliceBy(o2::framework::PresliceBase<T1, OPT, SORTED> const& container, int value) const
2102+
template <typename T1, typename Policy, bool OPT>
2103+
auto sliceBy(o2::framework::PresliceBase<T1, Policy, OPT> const& container, int value) const
21042104
{
21052105
return doSliceBy(this, container, value);
21062106
}
@@ -3201,8 +3201,8 @@ struct JoinFull : Table<o2::aod::Hash<"JOIN"_h>, D, o2::aod::Hash<"JOIN"_h>, Ts.
32013201
return doSliceByCachedUnsorted(this, node, value, cache);
32023202
}
32033203

3204-
template <typename T1, bool OPT, bool SORTED>
3205-
auto sliceBy(o2::framework::PresliceBase<T1, OPT, SORTED> const& container, int value) const
3204+
template <typename T1, typename Policy, bool OPT>
3205+
auto sliceBy(o2::framework::PresliceBase<T1, Policy, OPT> const& container, int value) const
32063206
{
32073207
return doSliceBy(this, container, value);
32083208
}
@@ -3463,14 +3463,16 @@ class FilteredBase : public T
34633463
return doSliceByCachedUnsorted(this, node, value, cache);
34643464
}
34653465

3466-
template <typename T1, bool OPT, bool SORTED>
3467-
auto sliceBy(o2::framework::PresliceBase<T1, OPT, SORTED> const& container, int value) const
3466+
template <typename T1, bool OPT>
3467+
auto sliceBy(o2::framework::PresliceBase<T1, framework::PreslicePolicySorted, OPT> const& container, int value) const
34683468
{
3469-
if constexpr (SORTED) {
3470-
return doFilteredSliceBy(this, container, value);
3471-
} else {
3472-
return doSliceBy(this, container, value);
3473-
}
3469+
return doFilteredSliceBy(this, container, value);
3470+
}
3471+
3472+
template <typename T1, bool OPT>
3473+
auto sliceBy(o2::framework::PresliceBase<T1, framework::PreslicePolicyGeneral, OPT> const& container, int value) const
3474+
{
3475+
return doSliceBy(this, container, value);
34743476
}
34753477

34763478
auto select(framework::expressions::Filter const& f) const
@@ -3697,14 +3699,16 @@ class Filtered : public FilteredBase<T>
36973699
return doSliceByCachedUnsorted(this, node, value, cache);
36983700
}
36993701

3700-
template <typename T1, bool OPT, bool SORTED>
3701-
auto sliceBy(o2::framework::PresliceBase<T1, OPT, SORTED> const& container, int value) const
3702+
template <typename T1, bool OPT>
3703+
auto sliceBy(o2::framework::PresliceBase<T1, framework::PreslicePolicySorted, OPT> const& container, int value) const
37023704
{
3703-
if constexpr (SORTED) {
3704-
return doFilteredSliceBy(this, container, value);
3705-
} else {
3706-
return doSliceBy(this, container, value);
3707-
}
3705+
return doFilteredSliceBy(this, container, value);
3706+
}
3707+
3708+
template <typename T1, bool OPT>
3709+
auto sliceBy(o2::framework::PresliceBase<T1, framework::PreslicePolicyGeneral, OPT> const& container, int value) const
3710+
{
3711+
return doSliceBy(this, container, value);
37083712
}
37093713

37103714
auto select(framework::expressions::Filter const& f) const
@@ -3864,14 +3868,16 @@ class Filtered<Filtered<T>> : public FilteredBase<typename T::table_t>
38643868
return doSliceByCachedUnsorted(this, node, value, cache);
38653869
}
38663870

3867-
template <typename T1, bool OPT, bool SORTED>
3868-
auto sliceBy(o2::framework::PresliceBase<T1, OPT, SORTED> const& container, int value) const
3871+
template <typename T1, bool OPT>
3872+
auto sliceBy(o2::framework::PresliceBase<T1, framework::PreslicePolicySorted, OPT> const& container, int value) const
38693873
{
3870-
if constexpr (SORTED) {
3871-
return doFilteredSliceBy(this, container, value);
3872-
} else {
3873-
return doSliceBy(this, container, value);
3874-
}
3874+
return doFilteredSliceBy(this, container, value);
3875+
}
3876+
3877+
template <typename T1, bool OPT>
3878+
auto sliceBy(o2::framework::PresliceBase<T1, framework::PreslicePolicyGeneral, OPT> const& container, int value) const
3879+
{
3880+
return doSliceBy(this, container, value);
38753881
}
38763882

38773883
private:

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -652,8 +652,8 @@ struct Partition {
652652
return mFiltered->sliceByCachedUnsorted(node, value, cache);
653653
}
654654

655-
template <typename T1, bool OPT, bool SORTED>
656-
[[nodiscard]] auto sliceBy(o2::framework::PresliceBase<T1, OPT, SORTED> const& container, int value) const
655+
template <typename T1, typename Policy, bool OPT>
656+
[[nodiscard]] auto sliceBy(o2::framework::PresliceBase<T1, Policy, OPT> const& container, int value) const
657657
{
658658
return mFiltered->sliceBy(container, value);
659659
}

0 commit comments

Comments
 (0)