Skip to content

Commit 6e8bcc6

Browse files
aalkinaalkin
committed
join/concat using ranges
1 parent c0b5e59 commit 6e8bcc6

3 files changed

Lines changed: 100 additions & 66 deletions

File tree

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

Lines changed: 16 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1238,7 +1238,6 @@ struct TableIterator : IP, C... {
12381238
{
12391239
static_assert(std::same_as<decltype(&(static_cast<B*>(this)->mColumnIterator)), std::decay_t<decltype(B::mColumnIterator)>*>, "foo");
12401240
return &(static_cast<B*>(this)->mColumnIterator);
1241-
// return static_cast<std::decay_t<decltype(B::mColumnIterator)>*>(nullptr);
12421241
}
12431242

12441243
template <typename B>
@@ -1249,10 +1248,10 @@ struct TableIterator : IP, C... {
12491248
};
12501249

12511250
struct ArrowHelpers {
1252-
static std::shared_ptr<arrow::Table> joinTables(std::vector<std::shared_ptr<arrow::Table>>&& tables);
1253-
static std::shared_ptr<arrow::Table> joinTables(std::vector<std::shared_ptr<arrow::Table>>&& tables, std::span<const char* const> labels);
1254-
static std::shared_ptr<arrow::Table> joinTables(std::vector<std::shared_ptr<arrow::Table>>&& tables, std::span<const std::string> labels);
1255-
static std::shared_ptr<arrow::Table> concatTables(std::vector<std::shared_ptr<arrow::Table>>&& tables);
1251+
static o2::soa::ArrowTableRef joinTables(std::vector<o2::soa::ArrowTableRef>&& tables);
1252+
static o2::soa::ArrowTableRef joinTables(std::vector<o2::soa::ArrowTableRef>&& tables, std::span<const char* const> labels);
1253+
static o2::soa::ArrowTableRef joinTables(std::vector<o2::soa::ArrowTableRef>&& tables, std::span<const std::string> labels);
1254+
static o2::soa::ArrowTableRef concatTables(std::vector<o2::soa::ArrowTableRef>&& tables);
12561255
};
12571256

12581257
template <size_t N1, std::array<TableRef, N1> os1, size_t N2, std::array<TableRef, N2> os2>
@@ -1963,26 +1962,6 @@ class Table
19631962
}
19641963
}
19651964

1966-
// Table(std::shared_ptr<arrow::Table> table, uint64_t offset = 0)
1967-
// : mTable(table),
1968-
// mOffset(offset),
1969-
// mEnd{table->num_rows()}
1970-
// {
1971-
// if (mTable->num_rows() == 0) {
1972-
// for (size_t ci = 0; ci < framework::pack_size(columns_t{}); ++ci) {
1973-
// mColumnChunks[ci] = nullptr;
1974-
// }
1975-
// mBegin = mEnd;
1976-
// } else {
1977-
// auto lookups = [this]<typename... C>(framework::pack<C...>) -> std::array<arrow::ChunkedArray*, framework::pack_size(columns_t{})> { return {lookupColumn<C>()...}; }(columns_t{});
1978-
// for (size_t ci = 0; ci < framework::pack_size(columns_t{}); ++ci) {
1979-
// mColumnChunks[ci] = lookups[ci];
1980-
// }
1981-
// mBegin = unfiltered_iterator{mColumnChunks, {table->num_rows(), offset}};
1982-
// mBegin.bindInternalIndices(this);
1983-
// }
1984-
// }
1985-
19861965
Table(std::vector<std::shared_ptr<arrow::Table>>&& tables, ArrowRange range)
19871966
requires(ref.origin_hash != "CONC"_h)
19881967
: Table({ArrowHelpers::joinTables(std::move(tables), std::span{originalLabels}), range})
@@ -1995,6 +1974,18 @@ class Table
19951974
{
19961975
}
19971976

1977+
Table(std::vector<o2::soa::ArrowTableRef>&& tables)
1978+
requires(ref.origin_hash != "CONC"_h)
1979+
: Table(ArrowHelpers::joinTables(std::move(tables), std::span{originalLabels}))
1980+
{
1981+
}
1982+
1983+
Table(std::vector<o2::soa::ArrowTableRef>&& tables)
1984+
requires(ref.origin_hash == "CONC"_h)
1985+
: Table(ArrowHelpers::concatTables(std::move(tables)))
1986+
{
1987+
}
1988+
19981989
template <typename Key>
19991990
inline arrow::ChunkedArray* getIndexToKey()
20001991
{

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111

1212
#ifndef O2_FRAMEWORK_ARROWTYPES_H
1313
#define O2_FRAMEWORK_ARROWTYPES_H
14+
#include <arrow/table.h>
1415
#include "Framework/Traits.h"
1516
#include "arrow/type_fwd.h"
1617
#include <span>
@@ -20,12 +21,29 @@ namespace o2::soa
2021
struct ArrowRange {
2122
uint64_t offset;
2223
int64_t size;
24+
25+
bool operator!=(ArrowRange const& other) const
26+
{
27+
return (offset != other.offset) && (size != other.size);
28+
}
2329
};
2430

2531
struct ArrowTableRef {
2632
std::shared_ptr<arrow::Table> tablePtr = nullptr;
2733
ArrowRange range{0, 0};
2834

35+
ArrowTableRef() = default;
36+
ArrowTableRef(std::shared_ptr<arrow::Table> table)
37+
: tablePtr{table},
38+
range{0, table->num_rows()}
39+
{
40+
}
41+
ArrowTableRef(std::shared_ptr<arrow::Table> table, ArrowRange range_)
42+
: tablePtr{table},
43+
range{range_}
44+
{
45+
}
46+
2947
ArrowTableRef makeEmpty() const
3048
{
3149
return {tablePtr, {0, 0}};
@@ -35,6 +53,11 @@ struct ArrowTableRef {
3553
{
3654
return {tablePtr, newRange};
3755
}
56+
57+
std::shared_ptr<arrow::Table> const& operator->() const
58+
{
59+
return tablePtr;
60+
}
3861
};
3962

4063
template <typename T>

‎Framework/Core/src/ASoA.cxx‎

Lines changed: 61 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -69,21 +69,6 @@ SelectionVector sliceSelection(std::span<int64_t const> const& mSelectedRows, in
6969
return slicedSelection;
7070
}
7171

72-
std::shared_ptr<arrow::Table> ArrowHelpers::joinTables(std::vector<std::shared_ptr<arrow::Table>>&& tables)
73-
{
74-
std::vector<std::shared_ptr<arrow::Field>> fields;
75-
std::vector<std::shared_ptr<arrow::ChunkedArray>> columns;
76-
bool notEmpty = (tables[0]->num_rows() != 0);
77-
std::ranges::for_each(tables, [&fields, &columns, notEmpty](auto const& t) {
78-
std::ranges::copy(t->fields(), std::back_inserter(fields));
79-
if (notEmpty) {
80-
std::ranges::copy(t->columns(), std::back_inserter(columns));
81-
}
82-
});
83-
auto schema = std::make_shared<arrow::Schema>(fields);
84-
return arrow::Table::Make(schema, columns);
85-
}
86-
8772
namespace
8873
{
8974
template <typename T>
@@ -109,62 +94,97 @@ void canNotJoin(std::vector<std::shared_ptr<arrow::Table>> const& tables, std::s
10994
}
11095
}
11196
}
112-
} // namespace
11397

114-
std::shared_ptr<arrow::Table> ArrowHelpers::joinTables(std::vector<std::shared_ptr<arrow::Table>>&& tables, std::span<const char* const> labels)
115-
{
116-
if (tables.size() == 1) {
117-
return tables[0];
98+
template <typename T>
99+
void IncompatibleRanges(std::vector<ArrowTableRef> const& tables, std::span<T> labels)
100+
{
101+
auto loc = std::ranges::adjacent_find(tables, [](auto const& l, auto const& r){ return l.range != r.range; });
102+
if (loc != std::ranges::cend(tables)) {
103+
auto pos = std::distance(tables.begin(), loc);
104+
auto next = loc + 1;
105+
if (labels.empty()) {
106+
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);
107+
} else {
108+
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);
109+
}
118110
}
119-
canNotJoin(tables, labels);
120-
return joinTables(std::forward<std::vector<std::shared_ptr<arrow::Table>>>(tables));
121111
}
122112

123-
std::shared_ptr<arrow::Table> ArrowHelpers::joinTables(std::vector<std::shared_ptr<arrow::Table>>&& tables, std::span<const std::string> labels)
113+
std::shared_ptr<arrow::Table> joinTablesImpl(std::ranges::input_range auto tables)
114+
{
115+
std::vector<std::shared_ptr<arrow::Field>> fields;
116+
std::vector<std::shared_ptr<arrow::ChunkedArray>> columns;
117+
bool notEmpty = (tables.front()->num_rows() != 0);
118+
std::ranges::for_each(tables, [&fields, &columns, notEmpty](auto const& t) {
119+
std::ranges::copy(t->fields(), std::back_inserter(fields));
120+
if (notEmpty) {
121+
std::ranges::copy(t->columns(), std::back_inserter(columns));
122+
}
123+
});
124+
auto schema = std::make_shared<arrow::Schema>(fields);
125+
return arrow::Table::Make(schema, columns);
126+
}
127+
128+
template <typename T>
129+
ArrowTableRef joinTablesImpl(std::ranges::input_range auto tables, std::span<T> labels)
124130
{
125131
if (tables.size() == 1) {
126-
return tables[0];
132+
return tables.front();
127133
}
128-
canNotJoin(tables, labels);
129-
return joinTables(std::forward<std::vector<std::shared_ptr<arrow::Table>>>(tables));
134+
IncompatibleRanges(tables, labels);
135+
ArrowRange commonRange{tables.front().range};
136+
return {joinTablesImpl(tables), commonRange};
137+
}
138+
} // namespace
139+
140+
o2::soa::ArrowTableRef ArrowHelpers::joinTables(std::vector<o2::soa::ArrowTableRef>&& tables)
141+
{
142+
return joinTablesImpl(tables, std::span<const char* const>());
130143
}
131144

132-
std::shared_ptr<arrow::Table> ArrowHelpers::concatTables(std::vector<std::shared_ptr<arrow::Table>>&& tables)
145+
o2::soa::ArrowTableRef ArrowHelpers::joinTables(std::vector<o2::soa::ArrowTableRef>&& tables, std::span<const char* const> labels)
146+
{
147+
return joinTablesImpl(tables, labels);
148+
}
149+
150+
o2::soa::ArrowTableRef ArrowHelpers::joinTables(std::vector<o2::soa::ArrowTableRef>&& tables, std::span<const std::string> labels)
151+
{
152+
return joinTablesImpl(tables, labels);
153+
}
154+
155+
o2::soa::ArrowTableRef ArrowHelpers::concatTables(std::vector<o2::soa::ArrowTableRef>&& tables)
133156
{
134157
if (tables.size() == 1) {
135-
return tables[0];
158+
return tables.front();
136159
}
137160
std::vector<std::shared_ptr<arrow::ChunkedArray>> columns;
138161
std::vector<std::shared_ptr<arrow::Field>> resultFields = tables[0]->schema()->fields();
139162
auto compareFields = [](std::shared_ptr<arrow::Field> const& f1, std::shared_ptr<arrow::Field> const& f2) {
140163
// Let's do this with stable sorting.
141164
return (!f1->Equals(f2)) && (f1->name() < f2->name());
142165
};
143-
for (size_t i = 1; i < tables.size(); ++i) {
144-
auto& fields = tables[i]->schema()->fields();
166+
std::ranges::for_each(tables.begin() + 1, tables.end(), [&resultFields, &compareFields](auto const& ref) mutable {
167+
std::vector<std::shared_ptr<arrow::Field>> const& fields = ref->fields();
145168
std::vector<std::shared_ptr<arrow::Field>> intersection;
146-
147-
std::set_intersection(resultFields.begin(), resultFields.end(),
148-
fields.begin(), fields.end(),
149-
std::back_inserter(intersection), compareFields);
169+
std::ranges::set_intersection(resultFields, fields, std::back_inserter(intersection), compareFields);
150170
resultFields.swap(intersection);
151-
}
171+
});
152172

153-
for (auto& field : resultFields) {
173+
std::ranges::transform(resultFields, std::back_inserter(columns), [&tables](auto const& field){
154174
arrow::ArrayVector chunks;
155-
for (auto& table : tables) {
175+
std::ranges::for_each(tables, [&field, &chunks](auto const& table){
156176
auto ci = table->schema()->GetFieldIndex(field->name());
157177
if (ci == -1) {
158178
throw std::runtime_error("Unable to find field " + field->name());
159179
}
160180
auto column = table->column(ci);
161181
auto otherChunks = column->chunks();
162182
chunks.insert(chunks.end(), otherChunks.begin(), otherChunks.end());
163-
}
164-
columns.push_back(std::make_shared<arrow::ChunkedArray>(chunks));
165-
}
183+
});
184+
return std::make_shared<arrow::ChunkedArray>(chunks);
185+
});
166186

167-
return arrow::Table::Make(std::make_shared<arrow::Schema>(resultFields), columns);
187+
return {arrow::Table::Make(std::make_shared<arrow::Schema>(resultFields), columns)};
168188
}
169189

170190
// ASCII-only lowercase. Column labels are plain identifiers, so we deliberately

0 commit comments

Comments
 (0)