Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions Framework/Core/include/Framework/ArrowTypes.h
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
#ifndef O2_FRAMEWORK_ARROWTYPES_H
#define O2_FRAMEWORK_ARROWTYPES_H
#include "arrow/type_fwd.h"
#include <span>

namespace o2::soa
{
Expand Down Expand Up @@ -62,6 +63,10 @@ template <>
struct arrow_array_for<double> {
using type = arrow::DoubleArray;
};
template <>
struct arrow_array_for<std::span<std::byte>> {
using type = arrow::BinaryViewArray;
};
template <int N>
struct arrow_array_for<float[N]> {
using type = arrow::FixedSizeListArray;
Expand Down
33 changes: 32 additions & 1 deletion Framework/Core/include/Framework/TableBuilder.h
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,7 @@ O2_ARROW_STL_CONVERSION(long unsigned, UInt64Type)
O2_ARROW_STL_CONVERSION(float, FloatType)
O2_ARROW_STL_CONVERSION(double, DoubleType)
O2_ARROW_STL_CONVERSION(std::string, StringType)
O2_ARROW_STL_CONVERSION(std::span<std::byte>, BinaryViewType)
} // namespace detail

void addLabelToSchema(std::shared_ptr<arrow::Schema>& schema, const char* label);
Expand Down Expand Up @@ -274,6 +275,29 @@ struct BuilderMaker<bool> {
}
};

template <>
struct BuilderMaker<std::span<std::byte>> {
using FillType = std::span<std::byte>;
using STLValueType = std::span<std::byte>;
using ArrowType = typename detail::ConversionTraits<std::span<std::byte>>::ArrowType;
using BuilderType = typename arrow::TypeTraits<ArrowType>::BuilderType;

static std::unique_ptr<BuilderType> make(arrow::MemoryPool* pool)
{
return std::make_unique<BuilderType>(pool);
}

static std::shared_ptr<arrow::DataType> make_datatype()
{
return arrow::TypeTraits<ArrowType>::type_singleton();
}

static arrow::Status append(BuilderType& builder, std::span<std::byte> value)
{
return builder.Append((char*)value.data(), (int64_t)value.size());
}
};

template <typename ITERATOR>
struct BuilderMaker<std::pair<ITERATOR, ITERATOR>> {
using FillType = std::pair<ITERATOR, ITERATOR>;
Expand Down Expand Up @@ -422,6 +446,13 @@ struct DirectInsertion {
return builder->Append(value);
}

template <typename BUILDER>
requires std::same_as<std::span<std::byte>, T>
arrow::Status append(BUILDER& builder, T value)
{
return builder->Append((char*)value.data(), (int64_t)value.size());
}

template <typename BUILDER>
arrow::Status flush(BUILDER&)
{
Expand Down Expand Up @@ -569,7 +600,7 @@ template <typename... ARGS>
using IndexedHoldersTuple = decltype(makeHolderTypes<ARGS...>());

template <typename T>
concept ShouldNotDeconstruct = std::is_bounded_array_v<T> || std::is_arithmetic_v<T> || framework::is_base_of_template_v<std::vector, T>;
concept ShouldNotDeconstruct = std::is_bounded_array_v<T> || std::is_arithmetic_v<T> || framework::is_base_of_template_v<std::vector, T> || std::same_as<std::span<std::byte>, T>;

/// Helper class which creates a lambda suitable for building
/// an arrow table from a tuple. This can be used, for example
Expand Down
31 changes: 31 additions & 0 deletions Framework/Core/test/test_TableBuilder.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,9 @@
#include <arrow/ipc/writer.h>
#include <arrow/ipc/reader.h>

#include <iostream>
#include <string_view>

using namespace o2::framework;

// We use a different namespace to avoid clashes with the
Expand All @@ -27,10 +30,12 @@ namespace test2
{
DECLARE_SOA_COLUMN_FULL(X, x, uint64_t, "x");
DECLARE_SOA_COLUMN_FULL(Y, y, uint64_t, "y");
DECLARE_SOA_COLUMN_FULL(Blob, blob, std::span<std::byte>, "blob");
DECLARE_SOA_COLUMN_FULL(Pos, pos, int[4], "pos");
} // namespace test2

using TestTable = o2::soa::InPlaceTable<0, test2::X, test2::Y>;
using SpanTable = o2::soa::InPlaceTable<0, test2::Blob>;
using ArrayTable = o2::soa::InPlaceTable<0, test2::Pos>;

TEST_CASE("TestTableBuilder")
Expand Down Expand Up @@ -189,6 +194,32 @@ TEST_CASE("TestTableBuilderMore")
REQUIRE(table->schema()->field(3)->type()->id() == arrow::boolean()->id());
}

TEST_CASE("TestSpan")
{
TableBuilder builder;
std::vector<std::byte> buffer{10, std::byte{'c'}};
std::vector<std::byte> buffer1{10, std::byte{'a'}};

auto rowWriter = builder.persist<std::span<std::byte>>({"blob"});
rowWriter(0, std::span(buffer));
rowWriter(0, std::span(buffer));
rowWriter(0, std::span(buffer1.data(), 3));
rowWriter(0, std::span(buffer1.data(), 1));
auto table = builder.finalize();

REQUIRE(table->num_columns() == 1);
REQUIRE(table->num_rows() == 4);
REQUIRE(table->schema()->field(0)->name() == "blob");
REQUIRE(table->schema()->field(0)->type()->id() == arrow::binary_view()->id());

std::cout << table->ToString() << std::endl;

auto readBack = SpanTable{table};
for (auto& row : readBack) {
std::cout << row.blob().size() << std::endl;
}
}

TEST_CASE("TestSoAIntegration")
{
TableBuilder builder;
Expand Down