diff --git a/Framework/Core/include/Framework/ASoA.h b/Framework/Core/include/Framework/ASoA.h index ba6d58df9a639..9e9b1d40f4008 100644 --- a/Framework/Core/include/Framework/ASoA.h +++ b/Framework/Core/include/Framework/ASoA.h @@ -58,6 +58,14 @@ void missingFilterDeclaration(int hash, int ai); void notBoundTable(const char* tableName); void* extractCCDBPayload(char* payload, size_t size, TClass const* cl, const char* what); +// ASCII-only lowercase. Column labels are plain identifiers, so we deliberately +// avoid the locale-aware std::tolower: it goes through the C locale facet on +// every character and dominated getIndexFromLabel in profiles. +constexpr char asciiToLower(char c) +{ + return (c >= 'A' && c <= 'Z') ? static_cast(c + 32) : c; +} + template auto createFieldsFromColumns(framework::pack) { @@ -494,13 +502,13 @@ class ColumnIterator : ChunkingPolicy : mColumn{column}, mCurrent{nullptr}, mCurrentPos{nullptr}, + mGlobalOffset{nullptr}, mLast{nullptr}, mFirstIndex{0}, - mCurrentChunk{0}, - mOffset{0} + mCurrentChunk{0} { auto array = getCurrentArray(); - mCurrent = reinterpret_cast const*>(array->values()->data()) + (mOffset >> SCALE_FACTOR); + mCurrent = reinterpret_cast const*>(array->values()->data()); mLast = mCurrent + array->length(); } @@ -516,10 +524,9 @@ class ColumnIterator : ChunkingPolicy { auto previousArray = getCurrentArray(); mFirstIndex += previousArray->length(); - mCurrentChunk++; auto array = getCurrentArray(); - mCurrent = reinterpret_cast const*>(array->values()->data()) + (mOffset >> SCALE_FACTOR) - (mFirstIndex >> SCALE_FACTOR); + mCurrent = reinterpret_cast const*>(array->values()->data()) - (mFirstIndex >> SCALE_FACTOR); mLast = mCurrent + array->length() + (mFirstIndex >> SCALE_FACTOR); } @@ -527,10 +534,9 @@ class ColumnIterator : ChunkingPolicy { auto previousArray = getCurrentArray(); mFirstIndex -= previousArray->length(); - mCurrentChunk--; auto array = getCurrentArray(); - mCurrent = reinterpret_cast const*>(array->values()->data()) + (mOffset >> SCALE_FACTOR) - (mFirstIndex >> SCALE_FACTOR); + mCurrent = reinterpret_cast const*>(array->values()->data()) - (mFirstIndex >> SCALE_FACTOR); mLast = mCurrent + array->length() + (mFirstIndex >> SCALE_FACTOR); } @@ -553,7 +559,7 @@ class ColumnIterator : ChunkingPolicy mCurrentChunk = mColumn->num_chunks() - 1; auto array = getCurrentArray(); mFirstIndex = mColumn->length() - array->length(); - mCurrent = reinterpret_cast const*>(array->values()->data()) + (mOffset >> SCALE_FACTOR) - (mFirstIndex >> SCALE_FACTOR); + mCurrent = reinterpret_cast const*>(array->values()->data()) - (mFirstIndex >> SCALE_FACTOR); mLast = mCurrent + array->length() + (mFirstIndex >> SCALE_FACTOR); } @@ -561,7 +567,7 @@ class ColumnIterator : ChunkingPolicy requires std::same_as> { checkSkipChunk(); - return (*(mCurrent - (mOffset >> SCALE_FACTOR) + ((*mCurrentPos + mOffset) >> SCALE_FACTOR)) & (1 << ((*mCurrentPos + mOffset) & 0x7))) != 0; + return (*(mCurrent + ((*mCurrentPos + *mGlobalOffset) >> SCALE_FACTOR)) & (1 << ((*mCurrentPos + *mGlobalOffset) & ((1 << SCALE_FACTOR) - 1)))) != 0; } auto operator*() const @@ -569,8 +575,8 @@ class ColumnIterator : ChunkingPolicy { checkSkipChunk(); auto list = std::static_pointer_cast(mColumn->chunk(mCurrentChunk)); - auto offset = list->value_offset(*mCurrentPos - mFirstIndex); - auto length = list->value_length(*mCurrentPos - mFirstIndex); + auto offset = list->value_offset(*mCurrentPos + *mGlobalOffset - mFirstIndex); + auto length = list->value_length(*mCurrentPos + *mGlobalOffset - mFirstIndex); return gsl::span const>{mCurrent + mFirstIndex + offset, mCurrent + mFirstIndex + (offset + length)}; } @@ -579,14 +585,14 @@ class ColumnIterator : ChunkingPolicy { checkSkipChunk(); auto array = std::static_pointer_cast(mColumn->chunk(mCurrentChunk)); - return array->GetView(*mCurrentPos - mFirstIndex); + return array->GetView(*mCurrentPos + *mGlobalOffset - mFirstIndex); } decltype(auto) operator*() const requires((!std::same_as>) && !std::same_as, arrow::ListArray> && !std::same_as, arrow::BinaryViewArray>) { checkSkipChunk(); - return *(mCurrent + (*mCurrentPos >> SCALE_FACTOR)); + return *(mCurrent + ((*mCurrentPos + *mGlobalOffset) >> SCALE_FACTOR)); } // Move to the chunk which containts element pos @@ -598,18 +604,18 @@ class ColumnIterator : ChunkingPolicy mutable unwrap_t const* mCurrent; int64_t const* mCurrentPos; + uint64_t const* mGlobalOffset; mutable unwrap_t const* mLast; arrow::ChunkedArray const* mColumn; mutable int mFirstIndex; mutable int mCurrentChunk; - mutable int mOffset; private: void checkSkipChunk() const requires((ChunkingPolicy::chunked == true) && std::same_as, arrow::ListArray>) { auto list = std::static_pointer_cast(mColumn->chunk(mCurrentChunk)); - if (O2_BUILTIN_UNLIKELY(*mCurrentPos - mFirstIndex >= list->length())) { + if (O2_BUILTIN_UNLIKELY(*mCurrentPos + *mGlobalOffset - mFirstIndex >= list->length())) { nextChunk(); } } @@ -617,7 +623,7 @@ class ColumnIterator : ChunkingPolicy void checkSkipChunk() const requires((ChunkingPolicy::chunked == true) && !std::same_as, arrow::ListArray>) { - if (O2_BUILTIN_UNLIKELY(((mCurrent + (*mCurrentPos >> SCALE_FACTOR)) >= mLast))) { + if (O2_BUILTIN_UNLIKELY(((mCurrent + ((*mCurrentPos + *mGlobalOffset) >> SCALE_FACTOR)) >= mLast))) { nextChunk(); } } @@ -631,7 +637,6 @@ class ColumnIterator : ChunkingPolicy requires(std::same_as, arrow::FixedSizeListArray>) { std::shared_ptr chunkToUse = mColumn->chunk(mCurrentChunk); - mOffset = chunkToUse->offset(); chunkToUse = std::dynamic_pointer_cast(chunkToUse)->values(); return std::static_pointer_cast>>(chunkToUse); } @@ -640,9 +645,7 @@ class ColumnIterator : ChunkingPolicy requires(std::same_as, arrow::ListArray>) { std::shared_ptr chunkToUse = mColumn->chunk(mCurrentChunk); - mOffset = chunkToUse->offset(); chunkToUse = std::dynamic_pointer_cast(chunkToUse)->values(); - mOffset = chunkToUse->offset(); return std::static_pointer_cast>>(chunkToUse); } @@ -650,7 +653,6 @@ class ColumnIterator : ChunkingPolicy requires(!std::same_as, arrow::FixedSizeListArray> && !std::same_as, arrow::ListArray>) { std::shared_ptr chunkToUse = mColumn->chunk(mCurrentChunk); - mOffset = chunkToUse->offset(); return std::static_pointer_cast>(chunkToUse); } }; @@ -1197,7 +1199,7 @@ struct TableIterator : IP, C... { { using namespace o2::soa; auto f = framework::overloaded{ - [this](T*) -> void { T::mColumnIterator.mCurrentPos = &this->mRowIndex; }, + [this](T*) -> void { T::mColumnIterator.mCurrentPos = &this->mRowIndex; T::mColumnIterator.mGlobalOffset = &this->mOffset; }, [this](T*) -> void { bindDynamicColumn(typename T::bindings_t{}); }, [this](T*) -> void {}, }; @@ -1225,7 +1227,6 @@ struct TableIterator : IP, C... { { static_assert(std::same_as(this)->mColumnIterator)), std::decay_t*>, "foo"); return &(static_cast(this)->mColumnIterator); - // return static_cast*>(nullptr); } template @@ -1236,10 +1237,14 @@ struct TableIterator : IP, C... { }; struct ArrowHelpers { - static std::shared_ptr joinTables(std::vector>&& tables); - static std::shared_ptr joinTables(std::vector>&& tables, std::span labels); - static std::shared_ptr joinTables(std::vector>&& tables, std::span labels); - static std::shared_ptr concatTables(std::vector>&& tables); + static o2::soa::ArrowTableRef joinTables(std::vector>&& tables); + static o2::soa::ArrowTableRef joinTables(std::vector&& tables); + static o2::soa::ArrowTableRef joinTables(std::vector&& tables, std::span labels); + static o2::soa::ArrowTableRef joinTables(std::vector&& tables, std::span labels); + static o2::soa::ArrowTableRef joinTables(std::vector>&& tables, std::span labels); + static o2::soa::ArrowTableRef joinTables(std::vector>&& tables, std::span labels); + static o2::soa::ArrowTableRef concatTables(std::vector&& tables); + static o2::soa::ArrowTableRef concatTables(std::vector>&& tables); }; template os1, size_t N2, std::array os2> @@ -1302,8 +1307,8 @@ static constexpr auto hasColumnForKey(framework::pack, std::string_view ke return std::ranges::equal( str1, str2, [](char c1, char c2) { - return std::tolower(static_cast(c1)) == - std::tolower(static_cast(c2)); + return asciiToLower(static_cast(c1)) == + asciiToLower(static_cast(c2)); }); }; return (caseInsensitiveCompare(C::inherited_t::mLabel, key) || ...); @@ -1403,12 +1408,7 @@ struct PreslicePolicySorted : public PreslicePolicyBase { void updateSliceInfo(SliceInfoPtr&& si); SliceInfoPtr sliceInfo; - std::shared_ptr getSliceFor(int value, std::shared_ptr const& input, uint64_t& offset) const; - // One-slot cache for the empty (0-row) slice, so that empty groups do not - // slice every column only to produce 0 rows (the common case for sparse - // grouping, e.g. candidates per collision). Keyed by the input table, which - // changes with every dataframe. - mutable std::pair> emptySlice{nullptr, nullptr}; + o2::soa::ArrowTableRef getSliceFor(int value, o2::soa::ArrowTableRef const& input) const; }; struct PreslicePolicyGeneral : public PreslicePolicyBase { @@ -1431,14 +1431,14 @@ struct PresliceBase : public Policy { { } - std::shared_ptr getSliceFor(int value, std::shared_ptr const& input, uint64_t& offset) const + o2::soa::ArrowTableRef getSliceFor(int value, o2::soa::ArrowTableRef const& input) const { if constexpr (OPT) { if (Policy::isMissing()) { - return nullptr; + return {nullptr, {0, 0}}; } } - return Policy::getSliceFor(value, input, offset); + return Policy::getSliceFor(value, input); } std::span getSliceFor(int value) const @@ -1507,9 +1507,8 @@ auto doSliceBy(T const* table, o2::framework::PresliceBase const missingOptionalPreslice(getLabelFromType>().data(), container.bindingKey.key.c_str()); } } - uint64_t offset = 0; - auto out = container.getSliceFor(value, table->asArrowTable(), offset); - auto t = typename T::self_t({out}, offset); + auto out = container.getSliceFor(value, table->asArrowTableRef()); + auto t = typename T::self_t({out}); if (t.tableSize() != 0) { table->copyIndexBindings(t); t.bindInternalIndicesTo(table); @@ -1520,7 +1519,7 @@ auto doSliceBy(T const* table, o2::framework::PresliceBase const template auto doSliceByHelper(T const* table, std::span const& selection) { - auto t = soa::Filtered({table->asArrowTable()}, selection); + auto t = soa::Filtered({table->asArrowTableRef()}, selection); if (t.tableSize() != 0) { table->copyIndexBindings(t); t.bindInternalIndicesTo(table); @@ -1533,7 +1532,7 @@ template requires(!soa::is_filtered_table) auto doSliceByHelper(T const* table, std::span const& selection) { - auto t = soa::Filtered({table->asArrowTable()}, selection); + auto t = soa::Filtered({table->asArrowTableRef()}, selection); if (t.tableSize() != 0) { table->copyIndexBindings(t); t.bindInternalIndicesTo(table); @@ -1557,17 +1556,17 @@ auto doSliceBy(T const* table, o2::framework::PresliceBase const SelectionVector sliceSelection(std::span const& mSelectedRows, int64_t nrows, uint64_t offset); template -auto prepareFilteredSlice(T const* table, std::shared_ptr slice, uint64_t offset) +auto prepareFilteredSlice(T const* table, o2::soa::ArrowTableRef slice) { - if (offset >= static_cast(table->tableSize())) { - Filtered fresult{{{slice}}, SelectionVector{}, 0}; + if (slice.range.offset >= static_cast(table->tableSize())) { + Filtered fresult{{slice}, SelectionVector{}}; if (fresult.tableSize() != 0) { table->copyIndexBindings(fresult); } return fresult; } - auto slicedSelection = sliceSelection(table->getSelectedRows(), slice->num_rows(), offset); - Filtered fresult{{{slice}}, std::move(slicedSelection), offset}; + auto slicedSelection = sliceSelection(table->getSelectedRows(), slice.range.size, slice.range.offset); + Filtered fresult{{slice}, std::move(slicedSelection)}; if (fresult.tableSize() != 0) { table->copyIndexBindings(fresult); } @@ -1583,9 +1582,8 @@ auto doFilteredSliceBy(T const* table, o2::framework::PresliceBase().data(), container.bindingKey.key.c_str()); } } - uint64_t offset = 0; - auto slice = container.getSliceFor(value, table->asArrowTable(), offset); - return prepareFilteredSlice(table, slice, offset); + auto slice = container.getSliceFor(value, table->asArrowTableRef()); + return prepareFilteredSlice(table, slice); } std::function originReplacement(header::DataOrigin newOrigin); @@ -1596,10 +1594,7 @@ auto doSliceByCached(T const* table, framework::expressions::BindingNode const& auto localCache = cache.ptr->getCacheFor({"", originReplacement(cache.ptr->newOrigin)(o2::soa::getMatcherFromTypeForKey(node.name)), node.name}); auto [offset, count] = localCache.getSliceFor(value); - // Empty group: reuse a cached empty (0-row) table instead of slicing every column. - auto slice = count == 0 ? cache.ptr->getEmptySliceFor(table->asArrowTable()) - : table->asArrowTable()->Slice(static_cast(offset), count); - auto t = typename T::self_t({slice}, static_cast(offset)); + auto t = typename T::self_t({table->asArrowTableRef().slice({static_cast(offset), count})}); if (t.tableSize() != 0) { table->copyIndexBindings(t); } @@ -1612,10 +1607,7 @@ auto doFilteredSliceByCached(T const* table, framework::expressions::BindingNode auto localCache = cache.ptr->getCacheFor({"", originReplacement(cache.ptr->newOrigin)(o2::soa::getMatcherFromTypeForKey(node.name)), node.name}); auto [offset, count] = localCache.getSliceFor(value); - // Empty group: reuse a cached empty (0-row) table instead of slicing every column. - auto slice = count == 0 ? cache.ptr->getEmptySliceFor(table->asArrowTable()) - : table->asArrowTable()->Slice(static_cast(offset), count); - return prepareFilteredSlice(table, slice, offset); + return prepareFilteredSlice(table, table->asArrowTableRef().slice({static_cast(offset), count})); } template @@ -1624,14 +1616,14 @@ auto doSliceByCachedUnsorted(T const* table, framework::expressions::BindingNode auto localCache = cache.ptr->getCacheUnsortedFor({"", originReplacement(cache.ptr->newOrigin)(o2::soa::getMatcherFromTypeForKey(node.name)), node.name}); if constexpr (soa::is_filtered_table) { - auto t = typename T::self_t({table->asArrowTable()}, localCache.getSliceFor(value)); + auto t = typename T::self_t({table->asArrowTableRef()}, localCache.getSliceFor(value)); if (t.tableSize() != 0) { t.intersectWithSelection(table->getSelectedRows()); table->copyIndexBindings(t); } return t; } else { - auto t = Filtered({table->asArrowTable()}, localCache.getSliceFor(value)); + auto t = Filtered({table->asArrowTableRef()}, localCache.getSliceFor(value)); if (t.tableSize() != 0) { table->copyIndexBindings(t); } @@ -1642,7 +1634,7 @@ auto doSliceByCachedUnsorted(T const* table, framework::expressions::BindingNode template auto select(T const& t, framework::expressions::Filter const& f) { - return Filtered({t.asArrowTable()}, selectionToVector(framework::expressions::createSelection(t.asArrowTable(), f))); + return Filtered({t.asArrowTableRef()}, selectionToVector(framework::expressions::createSelection(t.asArrowTable(), f))); } arrow::ChunkedArray* getIndexFromLabel(arrow::Table* table, std::string_view label); @@ -1651,7 +1643,6 @@ template consteval auto base_iter(framework::pack&&) -> TableIterator { } - template requires((sizeof...(Ts) > 0) && (soa::is_column && ...)) consteval auto getColumns() @@ -1709,7 +1700,7 @@ class Table static constexpr const auto originalLabels = [] refs, size_t... Is>(std::index_sequence) { return std::array{o2::aod::label()...}; }.template operator()(std::make_index_sequence()); - static constexpr const uint32_t binding_origin = originals[0].origin_hash; // commonOrigin(); + static constexpr const uint32_t binding_origin = originals[0].origin_hash; static constexpr header::DataOrigin binding_origin_ = o2::aod::Hash::origin; template bindings> @@ -1757,7 +1748,6 @@ class Table using columns_t = typename Parent::columns_t; using external_index_columns_t = typename Parent::external_index_columns_t; using bindings_pack_t = decltype([](framework::pack) -> framework::pack {}(external_index_columns_t{})); - // static constexpr const std::array originals{T::ref...}; static constexpr auto originals = Parent::originals; using policy_t = IP; using parent_t = Parent; @@ -1909,8 +1899,7 @@ class Table using iterator_template = TableIteratorBase; template - static consteval auto full_iter() - { + using iterator_template_o = decltype([]() { if constexpr (sizeof...(Ts) == 0) { return iterator_template{}; } else { @@ -1920,10 +1909,7 @@ class Table return iterator_template{}; } } - } - - template - using iterator_template_o = decltype(full_iter()); + }()); using iterator = iterator_template_o; using filtered_iterator = iterator_template_o; @@ -1937,12 +1923,11 @@ class Table return [](framework::pack) { return std::set{{C::hash...}}; }(columns_t{}); } - Table(std::shared_ptr table, uint64_t offset = 0) - : mTable(table), - mOffset(offset), - mEnd{table->num_rows()} + Table(o2::soa::ArrowTableRef tableRef) + : mArrowTableRef(tableRef), + mEnd{tableRef.range.size} { - if (mTable->num_rows() == 0) { + if (mArrowTableRef.tablePtr->num_rows() == 0) { for (size_t ci = 0; ci < framework::pack_size(columns_t{}); ++ci) { mColumnChunks[ci] = nullptr; } @@ -1952,20 +1937,37 @@ class Table for (size_t ci = 0; ci < framework::pack_size(columns_t{}); ++ci) { mColumnChunks[ci] = lookups[ci]; } - mBegin = unfiltered_iterator{mColumnChunks, {table->num_rows(), offset}}; + mBegin = unfiltered_iterator{mColumnChunks, {mEnd.index, mArrowTableRef.range.offset}}; mBegin.bindInternalIndices(this); } } - Table(std::vector>&& tables, uint64_t offset = 0) + Table(std::shared_ptr table) + : Table(o2::soa::ArrowTableRef{table}) + { + } + + Table(std::vector&& tables) + requires(ref.origin_hash != "CONC"_h) + : Table(ArrowHelpers::joinTables(std::forward>(tables), std::span{originalLabels})) + { + } + + Table(std::vector&& tables) + requires(ref.origin_hash == "CONC"_h) + : Table(ArrowHelpers::concatTables(std::forward>(tables))) + { + } + + Table(std::vector>&& tables) requires(ref.origin_hash != "CONC"_h) - : Table(ArrowHelpers::joinTables(std::move(tables), std::span{originalLabels}), offset) + : Table(ArrowHelpers::joinTables(std::forward>>(tables))) { } - Table(std::vector>&& tables, uint64_t offset = 0) + Table(std::vector>&& tables) requires(ref.origin_hash == "CONC"_h) - : Table(ArrowHelpers::concatTables(std::move(tables)), offset) + : Table(ArrowHelpers::concatTables(std::forward>>(tables))) { } @@ -2015,7 +2017,7 @@ class Table // is held by the table, so we are safe passing the bare pointer. If it does it // means that the iterator on a table is outliving the table itself, which is // a bad idea. - return filtered_iterator(mColumnChunks, {selection, mTable->num_rows(), mOffset}); + return filtered_iterator(mColumnChunks, {selection, mArrowTableRef.tablePtr->num_rows(), mArrowTableRef.range.offset}); } iterator iteratorAt(uint64_t i) const @@ -2043,17 +2045,27 @@ class Table /// Return a type erased arrow table backing store for / the type safe table. [[nodiscard]] std::shared_ptr asArrowTable() const { - return mTable; + return mArrowTableRef.tablePtr; + } + + [[nodiscard]] std::shared_ptr asArrowTableConstrained() const + { + return mArrowTableRef.tablePtr->Slice(mArrowTableRef.range.offset, mArrowTableRef.range.size); + } + + [[nodiscard]] ArrowTableRef asArrowTableRef() const + { + return mArrowTableRef; } /// Return offset auto offset() const { - return mOffset; + return mArrowTableRef.range.offset; } /// Size of the table, in rows. [[nodiscard]] int64_t size() const { - return mTable->num_rows(); + return mArrowTableRef.range.size; } [[nodiscard]] int64_t tableSize() const @@ -2139,27 +2151,30 @@ class Table auto rawSlice(uint64_t start, uint64_t end) const { - return self_t{mTable->Slice(start, end - start + 1), start}; + return self_t{mArrowTableRef.slice({start, static_cast(end - start + 1)})}; } auto emptySlice() const { - return self_t{mTable->Slice(0, 0), 0}; + return self_t{mArrowTableRef.makeEmpty()}; } private: template arrow::ChunkedArray* lookupColumn() { - if constexpr (soa::is_persistent_column) { - auto label = T::columnLabel(); - return getIndexFromLabel(mTable.get(), label); - } else { - return nullptr; - } + return nullptr; } - std::shared_ptr mTable = nullptr; - uint64_t mOffset = 0; + + template + arrow::ChunkedArray* lookupColumn() + { + return getIndexFromLabel(mArrowTableRef.tablePtr.get(), T::columnLabel()); + } + + ArrowTableRef mArrowTableRef; + // std::shared_ptr mTable = nullptr; + // uint64_t mOffset = 0; // Cached pointers to the ChunkedArray associated to a column arrow::ChunkedArray* mColumnChunks[framework::pack_size(columns_t{})]; RowViewSentinel mEnd; @@ -2262,13 +2277,13 @@ namespace o2::aod O2ORIGIN("AOD"); O2ORIGIN("AOD1"); O2ORIGIN("AOD2"); -// O2ORIGIN("DYN"); -// O2ORIGIN("IDX"); -// O2ORIGIN("ATIM"); + O2ORIGIN("JOIN"); O2HASH("JOIN/0"); + O2ORIGIN("CONC"); O2HASH("CONC/0"); + O2ORIGIN("TEST"); O2HASH("TEST/0"); } // namespace o2::aod @@ -3315,20 +3330,14 @@ struct Join : Table, o2::aod::Hash<"JOIN/0"_h>, o2::aod: static constexpr void isJoin() {}; using base = Table, o2::aod::Hash<"JOIN/0"_h>, o2::aod::Hash<"JOIN"_h>, Ts...>; - Join(std::shared_ptr&& table, uint64_t offset = 0) - : base{std::move(table), offset} - { - if (this->tableSize() != 0) { - bindInternalIndicesTo(this); - } - } - Join(std::vector>&& tables, uint64_t offset = 0) - : base{ArrowHelpers::joinTables(std::move(tables), std::span{base::originalLabels}), offset} + Join(std::vector&& tables) + : base{ArrowHelpers::joinTables(std::move(tables))} { if (this->tableSize() != 0) { bindInternalIndicesTo(this); } } + using base::bindExternalIndices; using base::bindInternalIndicesTo; static constexpr const uint32_t binding_origin = base::binding_origin; @@ -3398,12 +3407,12 @@ struct Join : Table, o2::aod::Hash<"JOIN/0"_h>, o2::aod: auto rawSlice(uint64_t start, uint64_t end) const { - return self_t{{this->asArrowTable()->Slice(start, end - start + 1)}, start}; + return self_t{{this->asArrowTableRef().slice({start, static_cast(end - start + 1)})}}; } auto emptySlice() const { - return self_t{{this->asArrowTable()->Slice(0, 0)}, 0}; + return self_t{{this->asArrowTableRef().slice({0, 0})}}; } template @@ -3418,7 +3427,7 @@ struct Join : Table, o2::aod::Hash<"JOIN/0"_h>, o2::aod: template constexpr auto join(Ts const&... t) { - return Join(ArrowHelpers::joinTables({t.asArrowTable()...}, std::span{Join::base::originalLabels})); + return Join({ArrowHelpers::joinTables({t.asArrowTableRef()...}, std::span{Join::base::originalLabels})}); } template @@ -3428,15 +3437,26 @@ template struct Concat : Table, o2::aod::Hash<"CONC/0"_h>, o2::aod::Hash<"CONC"_h>, Ts...> { using base = Table, o2::aod::Hash<"CONC/0"_h>, o2::aod::Hash<"CONC"_h>, Ts...>; using self_t = Concat; - Concat(std::vector>&& tables, uint64_t offset = 0) - : base{ArrowHelpers::concatTables(std::move(tables)), offset} + + Concat(ArrowTableRef table) + : base{table} { bindInternalIndicesTo(this); } - Concat(Ts const&... t, uint64_t offset = 0) - : base{ArrowHelpers::concatTables({t.asArrowTable()...}), offset} + + Concat(std::shared_ptr table) + : Concat{ArrowTableRef{table}} + { + } + + Concat(std::vector&& tables) + : Concat{ArrowHelpers::concatTables(std::move(tables))} + { + } + + Concat(Ts const&... t) + : Concat{ArrowHelpers::concatTables({t.asArrowTableRef()...})} { - bindInternalIndicesTo(this); } using base::originals; @@ -3462,6 +3482,9 @@ constexpr auto concat(Ts const&... t) return Concat{t...}; } +template +concept is_a_selection = std::same_as, gandiva::Selection> || std::same_as, SelectionVector> || std::same_as, std::span>; + template class FilteredBase : public T { @@ -3491,34 +3514,10 @@ class FilteredBase : public T using unfiltered_iterator = T::template iterator_template_o; using const_iterator = iterator; - FilteredBase(std::vector>&& tables, gandiva::Selection const& selection, uint64_t offset = 0) - : T{std::move(tables), offset}, - mSelectedRows{getSpan(selection)} - { - if (this->tableSize() != 0) { - mFilteredBegin = table_t::filtered_begin(mSelectedRows); - } - resetRanges(); - mFilteredBegin.bindInternalIndices(this); - } - - FilteredBase(std::vector>&& tables, SelectionVector&& selection, uint64_t offset = 0) - : T{std::move(tables), offset}, - mSelectedRowsCache{std::move(selection)}, - mCached{true} - { - mSelectedRows = std::span{mSelectedRowsCache}; - if (this->tableSize() != 0) { - mFilteredBegin = table_t::filtered_begin(mSelectedRows); - } - resetRanges(); - mFilteredBegin.bindInternalIndices(this); - } - - FilteredBase(std::vector>&& tables, std::span const& selection, uint64_t offset = 0) - : T{std::move(tables), offset}, - mSelectedRows{selection} + FilteredBase(std::vector&& tables, is_a_selection auto selection) + : T{std::move(tables)} { + adoptSelection(selection); if (this->tableSize() != 0) { mFilteredBegin = table_t::filtered_begin(mSelectedRows); } @@ -3570,7 +3569,7 @@ class FilteredBase : public T [[nodiscard]] int64_t tableSize() const { - return table_t::asArrowTable()->num_rows(); + return this->asArrowTableRef().range.size; } auto const& getSelectedRows() const @@ -3583,12 +3582,12 @@ class FilteredBase : public T SelectionVector newSelection; newSelection.resize(static_cast(end - start + 1)); std::iota(newSelection.begin(), newSelection.end(), start); - return self_t{{this->asArrowTable()}, std::move(newSelection), 0}; + return self_t{{this->asArrowTableRef()}, std::move(newSelection)}; } auto emptySlice() const { - return self_t{{this->asArrowTable()}, SelectionVector{}, 0}; + return self_t{{this->asArrowTableRef()}, SelectionVector{}}; } static inline auto getSpan(gandiva::Selection const& sel) @@ -3671,41 +3670,21 @@ class FilteredBase : public T return static_cast(std::distance(mSelectedRows.begin(), locate)); } - void sumWithSelection(SelectionVector const& selection) - { - mCached = true; - SelectionVector rowsUnion; - std::set_union(mSelectedRows.begin(), mSelectedRows.end(), selection.begin(), selection.end(), std::back_inserter(rowsUnion)); - mSelectedRowsCache.clear(); - mSelectedRowsCache = rowsUnion; - resetRanges(); - } - - void intersectWithSelection(SelectionVector const& selection) - { - mCached = true; - SelectionVector intersection; - std::set_intersection(mSelectedRows.begin(), mSelectedRows.end(), selection.begin(), selection.end(), std::back_inserter(intersection)); - mSelectedRowsCache.clear(); - mSelectedRowsCache = intersection; - resetRanges(); - } - - void sumWithSelection(std::span const& selection) + void sumWithSelection(is_a_selection auto selection) { mCached = true; SelectionVector rowsUnion; - std::set_union(mSelectedRows.begin(), mSelectedRows.end(), selection.begin(), selection.end(), std::back_inserter(rowsUnion)); + std::ranges::set_union(mSelectedRows, selection, std::back_inserter(rowsUnion)); mSelectedRowsCache.clear(); mSelectedRowsCache = rowsUnion; resetRanges(); } - void intersectWithSelection(std::span const& selection) + void intersectWithSelection(is_a_selection auto selection) { mCached = true; SelectionVector intersection; - std::set_intersection(mSelectedRows.begin(), mSelectedRows.end(), selection.begin(), selection.end(), std::back_inserter(intersection)); + std::ranges::set_intersection(mSelectedRows, selection, std::back_inserter(intersection)); mSelectedRowsCache.clear(); mSelectedRowsCache = intersection; resetRanges(); @@ -3730,6 +3709,36 @@ class FilteredBase : public T } } + template + inline void adoptSelection(S) + { + } + + template + requires(std::same_as, gandiva::Selection>) + inline void adoptSelection(S selection) + { + mSelectedRows = getSpan(selection); + mCached = false; + } + + template + requires(std::same_as, SelectionVector>) + inline void adoptSelection(S selection) + { + mSelectedRowsCache = std::move(selection); + mSelectedRows = std::span{mSelectedRowsCache}; + mCached = true; + } + + template + requires(std::same_as, std::span>) + inline void adoptSelection(S selection) + { + mSelectedRows = selection; + mCached = false; + } + std::span mSelectedRows; SelectionVector mSelectedRowsCache; bool mCached = false; @@ -3760,23 +3769,10 @@ class Filtered : public FilteredBase return const_iterator(this->cached_begin()); } - Filtered(std::vector>&& tables, gandiva::Selection const& selection, uint64_t offset = 0) - : FilteredBase(std::move(tables), selection, offset) {} - - Filtered(std::vector>&& tables, SelectionVector&& selection, uint64_t offset = 0) - : FilteredBase(std::move(tables), std::forward(selection), offset) {} + Filtered(std::vector&& tables, is_a_selection auto selection) + : FilteredBase{std::move(tables), std::forward(selection)} {} - Filtered(std::vector>&& tables, std::span const& selection, uint64_t offset = 0) - : FilteredBase(std::move(tables), selection, offset) {} - - Filtered operator+(SelectionVector const& selection) - { - Filtered copy(*this); - copy.sumWithSelection(selection); - return copy; - } - - Filtered operator+(std::span const& selection) + Filtered operator+(is_a_selection auto selection) { Filtered copy(*this); copy.sumWithSelection(selection); @@ -3788,13 +3784,7 @@ class Filtered : public FilteredBase return operator+(other.getSelectedRows()); } - Filtered operator+=(SelectionVector const& selection) - { - this->sumWithSelection(selection); - return *this; - } - - Filtered operator+=(std::span const& selection) + Filtered operator+=(is_a_selection auto selection) { this->sumWithSelection(selection); return *this; @@ -3805,14 +3795,7 @@ class Filtered : public FilteredBase return operator+=(other.getSelectedRows()); } - Filtered operator*(SelectionVector const& selection) - { - Filtered copy(*this); - copy.intersectWithSelection(selection); - return copy; - } - - Filtered operator*(std::span const& selection) + Filtered operator*(is_a_selection auto selection) { Filtered copy(*this); copy.intersectWithSelection(selection); @@ -3824,13 +3807,7 @@ class Filtered : public FilteredBase return operator*(other.getSelectedRows()); } - Filtered operator*=(SelectionVector const& selection) - { - this->intersectWithSelection(selection); - return *this; - } - - Filtered operator*=(std::span const& selection) + Filtered operator*=(is_a_selection auto selection) { this->intersectWithSelection(selection); return *this; @@ -3855,12 +3832,12 @@ class Filtered : public FilteredBase SelectionVector newSelection; newSelection.resize(static_cast(end - start + 1)); std::iota(newSelection.begin(), newSelection.end(), start); - return self_t{{this->asArrowTable()}, std::move(newSelection), 0}; + return self_t{{this->asArrowTableRef()}, std::move(newSelection)}; } auto emptySlice() const { - return self_t{{this->asArrowTable()}, SelectionVector{}, 0}; + return self_t{{this->asArrowTableRef()}, SelectionVector{}}; } template @@ -3922,38 +3899,15 @@ class Filtered> : public FilteredBase return const_iterator(this->cached_begin()); } - Filtered(std::vector>&& tables, gandiva::Selection const& selection, uint64_t offset = 0) - : FilteredBase(std::move(extractTablesFromFiltered(tables)), selection, offset) + Filtered(std::vector>&& tables, is_a_selection auto selection) + : FilteredBase(std::move(extractTablesFromFiltered(tables)), std::forward(selection)) { for (auto& table : tables) { *this *= table; } } - Filtered(std::vector>&& tables, SelectionVector&& selection, uint64_t offset = 0) - : FilteredBase(std::move(extractTablesFromFiltered(tables)), std::forward(selection), offset) - { - for (auto& table : tables) { - *this *= table; - } - } - - Filtered(std::vector>&& tables, std::span const& selection, uint64_t offset = 0) - : FilteredBase(std::move(extractTablesFromFiltered(tables)), selection, offset) - { - for (auto& table : tables) { - *this *= table; - } - } - - Filtered> operator+(SelectionVector const& selection) - { - Filtered> copy(*this); - copy.sumWithSelection(selection); - return copy; - } - - Filtered> operator+(std::span const& selection) + Filtered> operator+(is_a_selection auto selection) { Filtered> copy(*this); copy.sumWithSelection(selection); @@ -3965,13 +3919,7 @@ class Filtered> : public FilteredBase return operator+(other.getSelectedRows()); } - Filtered> operator+=(SelectionVector const& selection) - { - this->sumWithSelection(selection); - return *this; - } - - Filtered> operator+=(std::span const& selection) + Filtered> operator+=(is_a_selection auto selection) { this->sumWithSelection(selection); return *this; @@ -3982,14 +3930,7 @@ class Filtered> : public FilteredBase return operator+=(other.getSelectedRows()); } - Filtered> operator*(SelectionVector const& selection) - { - Filtered> copy(*this); - copy.intersectionWithSelection(selection); - return copy; - } - - Filtered> operator*(std::span const& selection) + Filtered> operator*(is_a_selection auto selection) { Filtered> copy(*this); copy.intersectionWithSelection(selection); @@ -4001,13 +3942,7 @@ class Filtered> : public FilteredBase return operator*(other.getSelectedRows()); } - Filtered> operator*=(SelectionVector const& selection) - { - this->intersectWithSelection(selection); - return *this; - } - - Filtered> operator*=(std::span const& selection) + Filtered> operator*=(is_a_selection auto selection) { this->intersectWithSelection(selection); return *this; @@ -4030,12 +3965,12 @@ class Filtered> : public FilteredBase SelectionVector newSelection; newSelection.resize(static_cast(end - start + 1)); std::iota(newSelection.begin(), newSelection.end(), start); - return self_t{{this->asArrowTable()}, std::move(newSelection), 0}; + return self_t{{this->asArrowTableRef()}, std::move(newSelection)}; } auto emptySlice() const { - return self_t{{this->asArrowTable()}, SelectionVector{}, 0}; + return self_t{{this->asArrowTableRef()}, SelectionVector{}}; } auto sliceByCached(framework::expressions::BindingNode const& node, int value, o2::framework::SliceCache& cache) const @@ -4061,11 +3996,11 @@ class Filtered> : public FilteredBase } private: - std::vector> extractTablesFromFiltered(std::vector>& tables) + std::vector extractTablesFromFiltered(std::vector>& tables) { - std::vector> outTables; + std::vector outTables; for (auto& table : tables) { - outTables.push_back(table.asArrowTable()); + outTables.push_back(table.asArrowTableRef()); } return outTables; } @@ -4101,15 +4036,13 @@ struct IndexTable : Table { ...); } - IndexTable(std::shared_ptr table, uint64_t offset = 0) - : base_t{table, offset} - { - } + IndexTable(ArrowTableRef table) + : base_t{table} {} - IndexTable(std::vector> tables, uint64_t offset = 0) - : base_t{tables[0], offset} - { - } + /// FIXME: this is a compatiblity for a generic constructor call with a vector + /// there has to be a safer way + IndexTable(std::vector&& tables) + : base_t{tables[0]} {} IndexTable(IndexTable const&) = default; IndexTable(IndexTable&&) = default; @@ -4126,14 +4059,9 @@ template struct SmallGroupsBase : public Filtered { static constexpr void isSmallGroups() {}; static constexpr bool applyFilters = APPLY; - SmallGroupsBase(std::vector>&& tables, gandiva::Selection const& selection, uint64_t offset = 0) - : Filtered(std::move(tables), selection, offset) {} - - SmallGroupsBase(std::vector>&& tables, SelectionVector&& selection, uint64_t offset = 0) - : Filtered(std::move(tables), std::forward(selection), offset) {} - SmallGroupsBase(std::vector>&& tables, std::span const& selection, uint64_t offset = 0) - : Filtered(std::move(tables), selection, offset) {} + SmallGroupsBase(std::vector&& tables, is_a_selection auto selection) + : Filtered(std::move(tables), selection) {} }; template diff --git a/Framework/Core/include/Framework/AnalysisHelpers.h b/Framework/Core/include/Framework/AnalysisHelpers.h index 96d89b722780b..e4d9682b35302 100644 --- a/Framework/Core/include/Framework/AnalysisHelpers.h +++ b/Framework/Core/include/Framework/AnalysisHelpers.h @@ -147,7 +147,7 @@ auto spawner(framework::pack, std::vector>&& if (fullTable->num_rows() == 0) { return makeEmptyTable(name, framework::pack{}); } - return spawnerHelper(fullTable, schema, sizeof...(C), projectors, name, projector); + return spawnerHelper(fullTable.tablePtr, schema, sizeof...(C), projectors, name, projector); } std::string serializeProjectors(std::vector& 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>>(std::vector{table.asArrowTable()}, std::forward(selection)); + return std::make_unique>>(std::vector{table.asArrowTableRef()}, std::forward(selection)); } void initializePartitionCaches(std::set const& hashes, std::shared_ptr 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 projectors{{std::move(Cs::Projector())...}}; static std::shared_ptr projector = nullptr; static auto schema = std::make_shared(o2::soa::createFieldsFromColumns(framework::pack{})); - return output_t{{o2::framework::spawner(framework::pack{}, {table.asArrowTable()}, "dynamicExtension", projectors.data(), projector, schema), table.asArrowTable()}, 0}; + return output_t{{o2::framework::spawner(framework::pack{}, {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 auto Attach(T const& table) { using output_t = Join, 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 diff --git a/Framework/Core/include/Framework/AnalysisManagers.h b/Framework/Core/include/Framework/AnalysisManagers.h index bb37fb9016c2f..8b426576ce556 100644 --- a/Framework/Core/include/Framework/AnalysisManagers.h +++ b/Framework/Core/include/Framework/AnalysisManagers.h @@ -319,12 +319,12 @@ bool prepareOutput(ProcessingContext& context, T& spawns) } using D = o2::aod::Hash; - spawns.extension = std::make_shared(o2::framework::spawner(originalTable, + spawns.extension = std::make_shared(o2::framework::spawner(originalTable.tablePtr, o2::aod::label(), spawns.projectors.data(), spawns.projector, spawns.schema)); - spawns.table = std::make_shared(soa::ArrowHelpers::joinTables({spawns.extension->asArrowTable(), originalTable}, std::span{T::spawnable_t::table_t::originalLabels})); + spawns.table = std::make_shared(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; - defines.extension = std::make_shared(o2::framework::spawner(originalTable, + defines.extension = std::make_shared(o2::framework::spawner(originalTable.tablePtr, o2::aod::label(), defines.projectors.data(), defines.projector, defines.schema)); - defines.table = std::make_shared(soa::ArrowHelpers::joinTables({defines.extension->asArrowTable(), originalTable}, std::span{T::spawnable_t::table_t::originalLabels})); + defines.table = std::make_shared(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; - defines.extension = std::make_shared(o2::framework::spawner(originalTable, + defines.extension = std::make_shared(o2::framework::spawner(originalTable.tablePtr, o2::aod::label(), defines.projectors.data(), defines.projector, defines.schema)); - defines.table = std::make_shared(soa::ArrowHelpers::joinTables({defines.extension->asArrowTable(), originalTable}, std::span{T::spawnable_t::table_t::originalLabels})); + defines.table = std::make_shared(soa::ArrowHelpers::joinTables({defines.extension->asArrowTableRef(), originalTable}, std::span{T::spawnable_t::table_t::originalLabels})); return true; } diff --git a/Framework/Core/include/Framework/AnalysisTask.h b/Framework/Core/include/Framework/AnalysisTask.h index 3170236e18f09..e4e6c35d9a7e9 100644 --- a/Framework/Core/include/Framework/AnalysisTask.h +++ b/Framework/Core/include/Framework/AnalysisTask.h @@ -226,7 +226,7 @@ struct AnalysisDataProcessorBuilder { template static auto extractTablesFromRecord(InputRecord& record, R matchers) { - std::vector> tables; + std::vector tables; std::ranges::transform(matchers, std::back_inserter(tables), [&record](auto const& m) { return record.get(m.second)->asArrowTable(); }); @@ -248,8 +248,8 @@ struct AnalysisDataProcessorBuilder { template static auto extractFilteredFromRecord(InputRecord& record, R matchers, ExpressionInfo& info) { - std::shared_ptr 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>) { if (info.selection == nullptr) { soa::missingFilterDeclaration(info.processHash, info.argumentIndex); diff --git a/Framework/Core/include/Framework/ArrowTypes.h b/Framework/Core/include/Framework/ArrowTypes.h index 2673472a81152..cd85db72315e2 100644 --- a/Framework/Core/include/Framework/ArrowTypes.h +++ b/Framework/Core/include/Framework/ArrowTypes.h @@ -11,12 +11,55 @@ #ifndef O2_FRAMEWORK_ARROWTYPES_H #define O2_FRAMEWORK_ARROWTYPES_H +#include #include "Framework/Traits.h" #include "arrow/type_fwd.h" #include 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 tablePtr = nullptr; + ArrowRange range{0, 0}; + + ArrowTableRef() = default; + ArrowTableRef(std::shared_ptr table) + : tablePtr{table}, + range{0, table->num_rows()} + { + } + ArrowTableRef(std::shared_ptr 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 const& operator->() const + { + return tablePtr; + } +}; + template struct arrow_array_for { }; diff --git a/Framework/Core/include/Framework/GroupSlicer.h b/Framework/Core/include/Framework/GroupSlicer.h index 74e5f16c1703f..f06cd6a0cd916 100644 --- a/Framework/Core/include/Framework/GroupSlicer.h +++ b/Framework/Core/include/Framework/GroupSlicer.h @@ -194,7 +194,7 @@ struct GroupSlicer { } } } - std::decay_t typedTable{{originalTable.asArrowTable()}, std::move(s)}; + std::decay_t 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{{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(offset); }); - std::decay_t typedTable{{groupedElementsTable}, std::move(slicedSelection), offset}; + std::decay_t typedTable{{groupedElementsTable}, std::move(slicedSelection)}; typedTable.bindInternalIndicesTo(&originalTable); return typedTable; } @@ -281,9 +272,6 @@ struct GroupSlicer { std::span groupSelection; std::array const*, sizeof...(A)> selections; std::array::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, sizeof...(A)> emptyTables{}; std::array sliceInfos; std::array sliceInfosUnsorted; diff --git a/Framework/Core/include/Framework/GroupedCombinations.h b/Framework/Core/include/Framework/GroupedCombinations.h index b0a6c9e658a10..d8c6aea44f31d 100644 --- a/Framework/Core/include/Framework/GroupedCombinations.h +++ b/Framework/Core/include/Framework/GroupedCombinations.h @@ -70,15 +70,15 @@ struct GroupedCombinationsGenerator { template GroupedIterator(const GroupingPolicy& groupingPolicy, const G& grouping, const std::tuple& associated, SliceCache* cache_) : GroupingPolicy(groupingPolicy), - mGrouping{std::make_shared(std::vector{grouping.asArrowTable()})}, + mGrouping{std::make_shared(std::vector{grouping.asArrowTableRef()})}, mAssociated{std::make_shared>(std::make_tuple(std::get(pack{})>(associated)...))}, mIndexColumns{getMatchingIndexNode()...}, cache{cache_} { if constexpr (soa::is_filtered_table>) { - mGrouping = std::make_shared(std::vector{grouping.asArrowTable()}, grouping.getSelectedRows()); + mGrouping = std::make_shared(std::vector{grouping.asArrowTableRef()}, grouping.getSelectedRows()); } else { - mGrouping = std::make_shared(std::vector{grouping.asArrowTable()}); + mGrouping = std::make_shared(std::vector{grouping.asArrowTableRef()}); } setMultipleGroupingTables(grouping); if (!this->mIsEnd) { @@ -94,9 +94,9 @@ struct GroupedCombinationsGenerator { void setTables(const G& grouping, const std::tuple& associated) { if constexpr (soa::is_filtered_table>) { - mGrouping = std::make_shared(std::vector{grouping.asArrowTable()}, grouping.getSelectedRows()); + mGrouping = std::make_shared(std::vector{grouping.asArrowTableRef()}, grouping.getSelectedRows()); } else { - mGrouping = std::make_shared(std::vector{grouping.asArrowTable()}); + mGrouping = std::make_shared(std::vector{grouping.asArrowTableRef()}); } mAssociated = std::make_shared>(std::make_tuple(std::get(pack{})>(associated)...)); setMultipleGroupingTables(grouping); diff --git a/Framework/Core/src/ASoA.cxx b/Framework/Core/src/ASoA.cxx index f565fa6e9ce47..486783c39d18a 100644 --- a/Framework/Core/src/ASoA.cxx +++ b/Framework/Core/src/ASoA.cxx @@ -69,21 +69,6 @@ SelectionVector sliceSelection(std::span const& mSelectedRows, in return slicedSelection; } -std::shared_ptr ArrowHelpers::joinTables(std::vector>&& tables) -{ - std::vector> fields; - std::vector> columns; - bool notEmpty = (tables[0]->num_rows() != 0); - std::ranges::for_each(tables, [&fields, &columns, notEmpty](auto const& t) { - std::ranges::copy(t->fields(), std::back_inserter(fields)); - if (notEmpty) { - std::ranges::copy(t->columns(), std::back_inserter(columns)); - } - }); - auto schema = std::make_shared(fields); - return arrow::Table::Make(schema, columns); -} - namespace { template @@ -109,53 +94,108 @@ void canNotJoin(std::vector> const& tables, std::s } } } -} // namespace -std::shared_ptr ArrowHelpers::joinTables(std::vector>&& tables, std::span labels) +template +void IncompatibleRanges(std::vector const& tables, std::span labels) +{ + auto loc = std::ranges::adjacent_find(tables, [](auto const& l, auto const& r) { return l.range != r.range; }); + if (loc != std::ranges::cend(tables)) { + auto pos = std::distance(tables.begin(), loc); + auto next = loc + 1; + if (labels.empty()) { + throw o2::framework::runtime_error_f("Incompatible ranges at %d: (%zu, %z) vs. (%zu, %z)", pos, loc->range.offset, loc->range.size, next->range.offset, next->range.size); + } else { + throw o2::framework::runtime_error_f("Incompatible ranges at %d between %s and %s: (%zu, %z) vs. (%zu, %z)", pos, makeString(labels[pos]), makeString(labels[pos + 1]), loc->range.offset, loc->range.size, next->range.offset, next->range.size); + } + } +} + +std::shared_ptr joinTablesImpl(std::ranges::input_range auto tables) +{ + std::vector> fields; + std::vector> columns; + bool notEmpty = (tables.front()->num_rows() != 0); + std::ranges::for_each(tables, [&fields, &columns, notEmpty](auto const& t) { + std::ranges::copy(t->fields(), std::back_inserter(fields)); + if (notEmpty) { + std::ranges::copy(t->columns(), std::back_inserter(columns)); + } + }); + auto schema = std::make_shared(fields); + return arrow::Table::Make(schema, columns); +} + +template +ArrowTableRef joinTablesImpl(std::ranges::input_range auto tables, std::span labels) { if (tables.size() == 1) { - return tables[0]; + return tables.front(); } + IncompatibleRanges(tables, labels); + ArrowRange commonRange{tables.front().range}; + return {joinTablesImpl(tables), commonRange}; +} +} // namespace + +o2::soa::ArrowTableRef ArrowHelpers::joinTables(std::vector>&& tables) +{ + std::vector refs; + std::ranges::transform(tables, std::back_inserter(refs), [](auto const& table) { return ArrowTableRef{table}; }); + return joinTablesImpl(refs, std::span()); +} + +o2::soa::ArrowTableRef ArrowHelpers::joinTables(std::vector&& tables) +{ + return joinTablesImpl(tables, std::span()); +} + +o2::soa::ArrowTableRef ArrowHelpers::joinTables(std::vector&& tables, std::span labels) +{ + return joinTablesImpl(tables, labels); +} + +o2::soa::ArrowTableRef ArrowHelpers::joinTables(std::vector&& tables, std::span labels) +{ + return joinTablesImpl(tables, labels); +} + +o2::soa::ArrowTableRef ArrowHelpers::joinTables(std::vector>&& tables, std::span labels) +{ canNotJoin(tables, labels); - return joinTables(std::forward>>(tables)); + return o2::soa::ArrowTableRef{joinTablesImpl(tables)}; } -std::shared_ptr ArrowHelpers::joinTables(std::vector>&& tables, std::span labels) +o2::soa::ArrowTableRef ArrowHelpers::joinTables(std::vector>&& tables, std::span labels) { - if (tables.size() == 1) { - return tables[0]; - } canNotJoin(tables, labels); - return joinTables(std::forward>>(tables)); + return o2::soa::ArrowTableRef{joinTablesImpl(tables)}; } -std::shared_ptr ArrowHelpers::concatTables(std::vector>&& tables) +o2::soa::ArrowTableRef ArrowHelpers::concatTables(std::vector&& tables) { if (tables.size() == 1) { - return tables[0]; + return tables.front(); } std::vector> columns; - std::vector> resultFields = tables[0]->schema()->fields(); + std::vector> resultFields = tables.front()->schema()->fields(); auto compareFields = [](std::shared_ptr const& f1, std::shared_ptr const& f2) { // Let's do this with stable sorting. return (!f1->Equals(f2)) && (f1->name() < f2->name()); }; - for (size_t i = 1; i < tables.size(); ++i) { - auto& fields = tables[i]->schema()->fields(); - std::vector> intersection; - std::set_intersection(resultFields.begin(), resultFields.end(), - fields.begin(), fields.end(), - std::back_inserter(intersection), compareFields); + for (auto i = 1; i < tables.size(); ++i) { + auto const& fields = tables[i]->fields(); + std::vector> intersection; + std::ranges::set_intersection(resultFields, fields, std::back_inserter(intersection), compareFields); resultFields.swap(intersection); } - for (auto& field : resultFields) { + for (auto const& field : resultFields) { arrow::ArrayVector chunks; - for (auto& table : tables) { + for (auto const& table : tables) { auto ci = table->schema()->GetFieldIndex(field->name()); if (ci == -1) { - throw std::runtime_error("Unable to find field " + field->name()); + throw framework::runtime_error_f("Unable to find field {}", field->name().c_str()); } auto column = table->column(ci); auto otherChunks = column->chunks(); @@ -164,15 +204,7 @@ std::shared_ptr ArrowHelpers::concatTables(std::vector(chunks)); } - return arrow::Table::Make(std::make_shared(resultFields), columns); -} - -// ASCII-only lowercase. Column labels are plain identifiers, so we deliberately -// avoid the locale-aware std::tolower: it goes through the C locale facet on -// every character and dominated getIndexFromLabel in profiles. -static constexpr char asciiToLower(char c) -{ - return (c >= 'A' && c <= 'Z') ? static_cast(c + 32) : c; + return {arrow::Table::Make(std::make_shared(resultFields), columns)}; } arrow::ChunkedArray* getIndexFromLabel(arrow::Table* table, std::string_view label) @@ -314,19 +346,10 @@ void PreslicePolicyGeneral::updateSliceInfo(SliceInfoUnsortedPtr&& si) sliceInfo = si; } -std::shared_ptr PreslicePolicySorted::getSliceFor(int value, std::shared_ptr const& input, uint64_t& offset) const +o2::soa::ArrowTableRef PreslicePolicySorted::getSliceFor(int value, o2::soa::ArrowTableRef const& input) const { auto [offset_, count] = this->sliceInfo.getSliceFor(value); - offset = static_cast(offset_); - if (count == 0) { - // Empty group: avoid slicing every column only to discard it. Cache one - // empty (0-row) table per input table and reuse it (see GroupSlicer). - if (emptySlice.first != input.get()) { - emptySlice = {input.get(), input->Slice(0, 0)}; - } - return emptySlice.second; - } - return input->Slice(offset_, count); + return input.slice({static_cast(offset_), count}); } std::span PreslicePolicyGeneral::getSliceFor(int value) const diff --git a/Framework/Core/src/AnalysisHelpers.cxx b/Framework/Core/src/AnalysisHelpers.cxx index 5e46ed86860e8..62229765ddbf8 100644 --- a/Framework/Core/src/AnalysisHelpers.cxx +++ b/Framework/Core/src/AnalysisHelpers.cxx @@ -207,7 +207,7 @@ std::shared_ptr Spawner::materialize(ProcessingContext& pc) const return arrow::Table::MakeEmpty(schema).ValueOrDie(); } - return spawnerHelper(fullTable, schema, binding.c_str(), schema->num_fields(), projector); + return spawnerHelper(fullTable.tablePtr, schema, binding.c_str(), schema->num_fields(), projector); } std::shared_ptr Builder::materialize(ProcessingContext& pc) diff --git a/Framework/Core/test/test_ASoA.cxx b/Framework/Core/test/test_ASoA.cxx index 48cfb277acc5c..2160fd884259c 100644 --- a/Framework/Core/test/test_ASoA.cxx +++ b/Framework/Core/test/test_ASoA.cxx @@ -106,7 +106,9 @@ TEST_CASE("TestTableIteration") auto i = ColumnIterator(table->column(0).get()); int64_t pos = 0; + uint64_t offset = 0; i.mCurrentPos = &pos; + i.mGlobalOffset = &offset; REQUIRE(*i == 0); pos++; REQUIRE(*i == 0); @@ -286,7 +288,7 @@ TEST_CASE("TestJoinedTables") REQUIRE(Test::contains()); REQUIRE(!Test::contains()); - Test tests{{tableX, tableY}, 0}; + Test tests{{tableX, tableY}}; REQUIRE(tests.contains()); REQUIRE(tests.contains()); @@ -308,7 +310,7 @@ TEST_CASE("TestJoinedTables") REQUIRE(15 == test.x() + test.y() + test.z()); } using TestMoreThanTwo = Join; - TestMoreThanTwo tests4{{tableX, tableY, tableZ}, 0}; + TestMoreThanTwo tests4{{tableX, tableY, tableZ}}; for (auto& test : tests4) { REQUIRE(15 == test.x() + test.y() + test.z()); } @@ -383,7 +385,7 @@ TEST_CASE("TestConcatTables") static_assert(std::same_as, o2::aod::test::Y, o2::aod::test::X, o2::aod::test::Z>>, "Bad nested join"); static_assert(std::same_as, o2::aod::test::X>>, "Bad intersection of columns"); - ConcatTest tests{tableA, tableB}; + ConcatTest tests{{tableA, tableB}}; REQUIRE(16 == tests.size()); for (auto& test : tests) { REQUIRE(test.index() == test.x()); @@ -428,7 +430,7 @@ TEST_CASE("TestConcatTables") gandiva::Selection selection_f = expressions::createSelection(tableA, testf); TestA testA{tableA}; - FilteredTest filtered{{testA.asArrowTable()}, selection_f}; + FilteredTest filtered{{testA.asArrowTableRef()}, selection_f}; REQUIRE(2 == filtered.size()); auto i = 0; @@ -451,7 +453,7 @@ TEST_CASE("TestConcatTables") selectionConcat->SetIndex(2, 10); selectionConcat->SetNumSlots(3); ConcatTest concatTest{tableA, tableB}; - FilteredConcatTest concatTestTable{{concatTest.asArrowTable()}, selectionConcat}; + FilteredConcatTest concatTestTable{{concatTest.asArrowTableRef()}, selectionConcat}; REQUIRE(3 == concatTestTable.size()); i = 0; @@ -480,8 +482,8 @@ TEST_CASE("TestConcatTables") selectionJoin->SetIndex(1, 2); selectionJoin->SetIndex(2, 4); selectionJoin->SetNumSlots(3); - JoinedTest testJoin{{tableA, tableC}, 0}; - FilteredJoinTest filteredJoin{{testJoin.asArrowTable()}, selectionJoin}; + JoinedTest testJoin{{tableA, tableC}}; + FilteredJoinTest filteredJoin{{testJoin.asArrowTableRef()}, selectionJoin}; i = 0; REQUIRE(filteredJoin.begin() != filteredJoin.end()); @@ -598,12 +600,12 @@ TEST_CASE("TestFilteredOperators") TestA testA{tableA}; auto s1 = expressions::createSelection(testA.asArrowTable(), f1); - FilteredTest filtered1{{testA.asArrowTable()}, s1}; + FilteredTest filtered1{{testA.asArrowTableRef()}, s1}; REQUIRE(4 == filtered1.size()); REQUIRE(filtered1.begin() != filtered1.end()); auto s2 = expressions::createSelection(testA.asArrowTable(), f2); - FilteredTest filtered2{{testA.asArrowTable()}, s2}; + FilteredTest filtered2{{testA.asArrowTableRef()}, s2}; REQUIRE(2 == filtered2.size()); REQUIRE(filtered2.begin() != filtered2.end()); @@ -631,7 +633,7 @@ TEST_CASE("TestFilteredOperators") expressions::Filter f3 = o2::aod::test::x < 3; auto s3 = expressions::createSelection(testA.asArrowTable(), f3); - FilteredTest filtered3{{testA.asArrowTable()}, s3}; + FilteredTest filtered3{{testA.asArrowTableRef()}, s3}; REQUIRE(3 == filtered3.size()); REQUIRE(filtered3.begin() != filtered3.end()); @@ -675,7 +677,7 @@ TEST_CASE("TestNestedFiltering") TestA testA{tableA}; auto s1 = expressions::createSelection(testA.asArrowTable(), f1); - FilteredTest filtered{{testA.asArrowTable()}, s1}; + FilteredTest filtered{{testA.asArrowTableRef()}, s1}; REQUIRE(4 == filtered.size()); REQUIRE(filtered.begin() != filtered.end()); @@ -718,7 +720,7 @@ TEST_CASE("TestEmptyTables") o2::aod::Infos i{iempty}; using PI = Join; - PI pi{{pempty, iempty}, 0}; + PI pi{{pempty, iempty}}; REQUIRE(pi.size() == 0); auto spawned = Extend(p); REQUIRE(spawned.size() == 0); @@ -772,7 +774,7 @@ TEST_CASE("TestIndexToFiltered") expressions::Filter flt = o2::aod::test::someBool == true; using Flt = o2::soa::Filtered; auto selection = expressions::createSelection(o.asArrowTable(), flt); - Flt f{{o.asArrowTable()}, selection}; + Flt f{{o.asArrowTableRef()}, selection}; r.bindExternalIndices(&f); auto it = r.begin(); it.moveByIndex(23); @@ -1084,7 +1086,7 @@ TEST_CASE("TestSelfIndexRecursion") } using FilteredPoints = o2::soa::Filtered; - FilteredPoints ffp({t1, t2}, {1, 2, 3}, 0); + FilteredPoints ffp({t1, t2}, SelectionVector{1, 2, 3}); ffp.bindInternalIndicesTo(&ffp); // Filter should not interfere with self-index and the binding should stay the same @@ -1252,6 +1254,63 @@ TEST_CASE("TestSliceByCachedMismatched") } } +TEST_CASE("TestSliceByCachedFiltered") +{ + TableBuilder b; + auto writer = b.cursor(); + for (auto i = 0; i < 20; ++i) { + writer(0, i, i % 3 == 0); + } + auto origins = b.finalize(); + o2::aod::Origints o{origins}; + + TableBuilder w; + auto writer_w = w.cursor(); + auto step = -1; + for (auto i = 0; i < 5 * 20; ++i) { + if (i % 5 == 0) { + ++step; + } + writer_w(0, step); + } + auto refs = w.finalize(); + o2::aod::References r{refs}; + + TableBuilder w2; + auto writer_w2 = w2.cursor(); + step = -1; + for (auto i = 0; i < 5 * 20; ++i) { + if (i % 3 == 0) { + ++step; + } + writer_w2(0, step); + } + auto refs2 = w2.finalize(); + o2::aod::OtherReferences r2{refs2}; + + using J = o2::soa::Join; + J rr{{refs, refs2}}; + + auto rrf = rr.select(o2::aod::test::altOrigintId > 2 && o2::aod::test::altOrigintId < 15); + + auto key = "fIndex" + o2::framework::cutString(o2::soa::getLabelFromType()) + "_alt"; + ArrowTableSlicingCache atscache({{o2::soa::getLabelFromTypeForKey(key), o2::soa::getMatcherFromTypeForKey(key), key}}); + auto s = atscache.updateCacheEntry(0, refs2); + SliceCache cache{&atscache}; + + for (auto& oi : o) { + auto cachedSlice = rrf.sliceByCached(o2::aod::test::altOrigintId, oi.globalIndex(), cache); + if (oi.globalIndex() <= 2 || oi.globalIndex() >= 15) { + CHECK(cachedSlice.size() == 0); + } else { + CHECK(cachedSlice.size() == 3); + } + for (auto& ri : cachedSlice) { + REQUIRE(ri.altOrigintId() == oi.globalIndex()); + } + } +} + TEST_CASE("TestIndexUnboundExceptions") { TableBuilder b; diff --git a/Framework/Core/test/test_ASoAHelpers.cxx b/Framework/Core/test/test_ASoAHelpers.cxx index c4d7f727aa295..701dc0bbced50 100644 --- a/Framework/Core/test/test_ASoAHelpers.cxx +++ b/Framework/Core/test/test_ASoAHelpers.cxx @@ -72,7 +72,7 @@ TEST_CASE("IteratorTuple") REQUIRE(*(static_cast(std::get<1>(maxOffset2)).getIterator().mCurrentPos) == 8); expressions::Filter filter = test::x > 3; - auto filtered = Filtered{{tests.asArrowTable()}, o2::framework::expressions::createSelection(tests.asArrowTable(), filter)}; + auto filtered = Filtered{{tests.asArrowTableRef()}, o2::framework::expressions::createSelection(tests.asArrowTable(), filter)}; std::tuple, Filtered> filteredTuple = std::make_tuple(filtered, filtered); auto it1 = std::get<0>(filteredTuple).begin(); @@ -164,7 +164,7 @@ TEST_CASE("CombinationsGeneratorConstruction") o2::framework::expressions::Filter filter = test::x > 3; auto s1 = o2::framework::expressions::createSelection(testsA.asArrowTable(), filter); - auto filtered = Filtered{{testsA.asArrowTable()}, s1}; + auto filtered = Filtered{{testsA.asArrowTableRef()}, s1}; CombinationsGenerator, Filtered>>::CombinationsIterator combItFiltered(CombinationsStrictlyUpperIndexPolicy(filtered, filtered)); REQUIRE(!(static_cast(std::get<0>(*(combItFiltered))).getIterator().mCurrentPos == nullptr)); diff --git a/Framework/Core/test/test_AnalysisDataModel.cxx b/Framework/Core/test/test_AnalysisDataModel.cxx index b8b9c161f0e07..ae0914a285110 100644 --- a/Framework/Core/test/test_AnalysisDataModel.cxx +++ b/Framework/Core/test/test_AnalysisDataModel.cxx @@ -49,9 +49,9 @@ TEST_CASE("TestJoinedTablesContains") using Test = o2::soa::Join; - Test tests{{tXY, tZD}, 0}; - REQUIRE(tests.asArrowTable()->num_columns() != 0); - REQUIRE(tests.asArrowTable()->num_columns() == + Test tests{{tXY, tZD}}; + REQUIRE(tests.asArrowTableRef()->num_columns() != 0); + REQUIRE(tests.asArrowTableRef()->num_columns() == tXY->num_columns() + tZD->num_columns()); auto tests2 = join(XY{tXY}, ZD{tZD}); static_assert(std::same_as, diff --git a/Framework/Core/test/test_AnalysisTask.cxx b/Framework/Core/test/test_AnalysisTask.cxx index f5d8c4c43bc38..cb710b9a3871c 100644 --- a/Framework/Core/test/test_AnalysisTask.cxx +++ b/Framework/Core/test/test_AnalysisTask.cxx @@ -314,7 +314,7 @@ TEST_CASE("TestPartitionIteration") expressions::Filter f1 = aod::test::x < 4.0f; auto selection = expressions::createSelection(testA.asArrowTable(), f1); - FilteredTest filtered{{testA.asArrowTable()}, o2::soa::selectionToVector(selection)}; + FilteredTest filtered{{testA.asArrowTableRef()}, o2::soa::selectionToVector(selection)}; PartitionFilteredTest p2 = aod::test::y > 9.0f; p2.bindTable(filtered); diff --git a/Framework/Core/test/test_Concepts.cxx b/Framework/Core/test/test_Concepts.cxx index ff5e0fa6200db..65703082519b6 100644 --- a/Framework/Core/test/test_Concepts.cxx +++ b/Framework/Core/test/test_Concepts.cxx @@ -121,7 +121,7 @@ TEST_CASE("IdentificationConcepts") REQUIRE(is_join); - auto tl = []() -> SmallGroups { return {std::vector>{}, SelectionVector{}, 0}; }; + auto tl = []() -> SmallGroups { return {{}, SelectionVector{}}; }; REQUIRE(is_smallgroups); // AnalysisHelpers diff --git a/Framework/Core/test/test_GroupSlicer.cxx b/Framework/Core/test/test_GroupSlicer.cxx index ee6878f23ff80..f282dcbc5c33b 100644 --- a/Framework/Core/test/test_GroupSlicer.cxx +++ b/Framework/Core/test/test_GroupSlicer.cxx @@ -195,8 +195,9 @@ TEST_CASE("GroupSlicerSeveralAssociated") {soa::getLabelFromType(), soa::getMatcherFromTypeForKey(key), key}, {soa::getLabelFromType(), soa::getMatcherFromTypeForKey(key), key}}); auto s = slices.updateCacheEntry(0, {trkTableX}); - s = slices.updateCacheEntry(1, {trkTableY}); - s = slices.updateCacheEntry(2, {trkTableZ}); + s &= slices.updateCacheEntry(1, {trkTableY}); + s &= slices.updateCacheEntry(2, {trkTableZ}); + REQUIRE(s.ok()); o2::framework::GroupSlicer g(e, tt, slices); auto count = 0; @@ -358,7 +359,7 @@ TEST_CASE("GroupSlicerMismatchedFilteredGroups") auto trkTable = builderT.finalize(); using FilteredEvents = soa::Filtered; soa::SelectionVector rows{2, 4, 10, 9, 15}; - FilteredEvents e{{evtTable}, {2, 4, 10, 9, 15}}; + FilteredEvents e{{{evtTable}}, soa::SelectionVector{2, 4, 10, 9, 15}}; aod::TrksX t{trkTable}; REQUIRE(e.size() == 5); REQUIRE(t.size() == 10 * (20 - 4)); @@ -419,7 +420,7 @@ TEST_CASE("GroupSlicerMismatchedUnsortedFilteredGroups") using FilteredEvents = soa::Filtered; soa::SelectionVector rows{2, 4, 10, 9, 15}; - FilteredEvents e{{evtTable}, {2, 4, 10, 9, 15}}; + FilteredEvents e{{evtTable}, soa::SelectionVector{2, 4, 10, 9, 15}}; soa::SmallGroups t{{trkTable}, std::move(sel)}; REQUIRE(e.size() == 5); @@ -631,9 +632,9 @@ TEST_CASE("EmptySliceables") TEST_CASE("ArrowDirectSlicing") { int counts[] = {5, 5, 5, 4, 1}; - int offsets[] = {0, 5, 10, 15, 19, 20}; + int const offsets[] = {0, 5, 10, 15, 19, 20}; int ids[] = {0, 1, 2, 3, 4}; - int sizes[] = {4, 1, 12, 5, 2}; + int const sizes[] = {4, 1, 12, 5, 2}; using BigE = soa::Join; @@ -683,34 +684,25 @@ TEST_CASE("ArrowDirectSlicing") REQUIRE(slices_vec[i]->length() == counts[i]); } - std::vector slices; - std::vector offsts; auto bk = Entry(soa::getLabelFromType(), soa::getMatcherFromTypeForKey("fID"), "fID"); ArrowTableSlicingCache cache({bk}); auto s = cache.updateCacheEntry(0, {evtTable}); + REQUIRE(s.ok()); auto lcache = cache.getCacheFor(bk); for (auto i = 0u; i < 5; ++i) { - auto [offset, count] = lcache.getSliceFor(i); - auto tbl = b_e.asArrowTable()->Slice(offset, count); - auto ca = tbl->GetColumnByName("fArr"); - auto cb = tbl->GetColumnByName("fBoo"); - auto cv = tbl->GetColumnByName("fLst"); - REQUIRE(ca->length() == counts[i]); - REQUIRE(cb->length() == counts[i]); - REQUIRE(cv->length() == counts[i]); - REQUIRE(ca->Equals(slices_array[i])); - REQUIRE(cb->Equals(slices_bool[i])); - REQUIRE(cv->Equals(slices_vec[i])); + auto [loffset, count] = lcache.getSliceFor(i); + auto tbl = b_e.asArrowTableRef().slice({static_cast(loffset), count}); + REQUIRE(tbl.range.size == counts[i]); } int j = 0u; for (auto i = 0u; i < 5; ++i) { - auto [offset, count] = lcache.getSliceFor(i); - auto tbl = BigE{{b_e.asArrowTable()->Slice(offset, count)}, static_cast(offset)}; + auto [loffset, count] = lcache.getSliceFor(i); + auto tbl = BigE{{b_e.asArrowTableRef().slice({static_cast(loffset), count})}}; REQUIRE(tbl.size() == counts[i]); for (auto& row : tbl) { REQUIRE(row.id() == ids[i]); - REQUIRE(row.boo() == (j % 2 == 0)); + CHECK(row.boo() == (j % 2 == 0)); auto rid = row.globalIndex(); auto arr = row.arr(); REQUIRE(arr[0] == 0.1f * (float)rid); @@ -729,7 +721,7 @@ TEST_CASE("ArrowDirectSlicing") TEST_CASE("TestSlicingException") { - int offsets[] = {0, 5, 10, 15, 19, 20}; + int const offsets[] = {0, 5, 10, 15, 19, 20}; int ids[] = {0, 1, 2, 4, 3}; TableBuilder builderE; diff --git a/Framework/Core/test/test_TableSpawner.cxx b/Framework/Core/test/test_TableSpawner.cxx index e200adf37ccb4..d5bd4c83068ac 100644 --- a/Framework/Core/test/test_TableSpawner.cxx +++ b/Framework/Core/test/test_TableSpawner.cxx @@ -53,7 +53,7 @@ TEST_CASE("TestTableSpawner") auto expoints_a = o2::soa::Extend(st1); Spawns s; auto extension = ExPointsExtension{o2::framework::spawner>(t1, o2::aod::Hash<"ExPoints"_h>::str, s.projectors.data(), s.projector, s.schema)}; - auto expoints = ExPoints{{t1, extension.asArrowTable()}, 0}; + auto expoints = ExPoints{{t1, extension.asArrowTable()}}; REQUIRE(expoints_a.size() == 9); REQUIRE(extension.size() == 9); @@ -81,7 +81,7 @@ TEST_CASE("TestTableSpawner") excpts.projectors[0] = test::x * test::x + test::y * test::y + test::z * test::z; auto extension_2 = ExcPointsCfgExtension{o2::framework::spawner>({t1}, o2::aod::Hash<"ExcPoints"_h>::str, excpts.projectors.data(), excpts.projector, excpts.schema)}; - auto excpoints = ExcPoints{{t1, extension_2.asArrowTable()}, 0}; + auto excpoints = ExcPoints{{t1, extension_2.asArrowTable()}}; rex = extension.begin(); auto rex_2 = extension_2.begin();