flatbuffer_utils.hpp

Go to the documentation of this file.

#pragma once
 
#include <ranges>
 
#include <flatbuffers/flatbuffers.h>
 
#include <sparrow/c_interface.hpp>
#include <sparrow/record_batch.hpp>
 
#include "File_generated.h"
#include "Message_generated.h"
 
#include "sparrow_ipc/compression.hpp"
#include "sparrow_ipc/utils.hpp"
 
namespace sparrow_ipc
{
    // Creates a Flatbuffers Decimal type from a format string
    // The format string is expected to be in the format "d:precision,scale"
    [[nodiscard]] std::pair<org::apache::arrow::flatbuf::Type, flatbuffers::Offset<void>>
    get_flatbuffer_decimal_type(
        flatbuffers::FlatBufferBuilder& builder,
        std::string_view format_str,
        const int32_t bitWidth
    );
 
    // Creates a Flatbuffers type from a format string
    // This function maps a sparrow data type to the corresponding Flatbuffers type
    [[nodiscard]] std::pair<org::apache::arrow::flatbuf::Type, flatbuffers::Offset<void>>
    get_flatbuffer_type(
        flatbuffers::FlatBufferBuilder& builder,
        std::string_view format_str,
        const std::optional<std::unordered_set<sparrow::ArrowFlag>>& flags = std::nullopt
    );

    [[nodiscard]] flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<org::apache::arrow::flatbuf::KeyValue>>>
    create_metadata(flatbuffers::FlatBufferBuilder& builder, const ArrowSchema& arrow_schema);

    [[nodiscard]] ::flatbuffers::Offset<org::apache::arrow::flatbuf::Field> create_field(
        flatbuffers::FlatBufferBuilder& builder,
        const ArrowSchema& arrow_schema,
        std::optional<std::string_view> name_override = std::nullopt,
        std::optional<int64_t> dictionary_id_override = std::nullopt
    );

    [[nodiscard]] ::flatbuffers::Offset<
        ::flatbuffers::Vector<::flatbuffers::Offset<org::apache::arrow::flatbuf::Field>>>
    create_children(flatbuffers::FlatBufferBuilder& builder, const sparrow::record_batch& record_batch);
 

    [[nodiscard]] ::flatbuffers::Offset<
        ::flatbuffers::Vector<::flatbuffers::Offset<org::apache::arrow::flatbuf::Field>>>
    create_children(flatbuffers::FlatBufferBuilder& builder, const ArrowSchema& arrow_schema);

    [[nodiscard]] flatbuffers::FlatBufferBuilder
    get_schema_message_builder(const sparrow::record_batch& record_batch);

    void fill_fieldnodes(
        const sparrow::arrow_proxy& arrow_proxy,
        std::vector<org::apache::arrow::flatbuf::FieldNode>& nodes
    );

    [[nodiscard]] std::vector<org::apache::arrow::flatbuf::FieldNode>
    create_fieldnodes(const sparrow::record_batch& record_batch);
 
    namespace details
    {
        std::size_t get_nb_buffers_to_process(const std::string_view& format, const std::size_t orig_buffers_size);
 
        template <typename Func>
        void fill_buffers_impl(
            const sparrow::arrow_proxy& arrow_proxy,
            std::vector<org::apache::arrow::flatbuf::Buffer>& flatbuf_buffers,
            int64_t& offset,
            Func&& get_buffer_size
        )
        {
            const auto& buffers = arrow_proxy.buffers();
            auto nb_buffers = get_nb_buffers_to_process(arrow_proxy.schema().format, buffers.size());
            std::ranges::for_each(buffers | std::views::take(nb_buffers),
                [&](const auto& buffer)
            {
                int64_t size = get_buffer_size(buffer);
                flatbuf_buffers.emplace_back(offset, size);
                offset += utils::align_to_8(size);
            });
 
            for (const auto& child : arrow_proxy.children())
            {
                fill_buffers_impl(child, flatbuf_buffers, offset, get_buffer_size);
            }
        }
 
        template <typename Func>
        std::vector<org::apache::arrow::flatbuf::Buffer>
        get_buffers_impl(const sparrow::record_batch& record_batch, Func&& fill_buffers_func)
        {
            std::vector<org::apache::arrow::flatbuf::Buffer> buffers;
            int64_t offset = 0;
            for (const auto& column : record_batch.columns())
            {
                const auto& arrow_proxy = sparrow::detail::array_access::get_arrow_proxy(column);
                fill_buffers_func(arrow_proxy, buffers, offset);
            }
            return buffers;
        }
    }  // namespace details

    void fill_buffers(
        const sparrow::arrow_proxy& arrow_proxy,
        std::vector<org::apache::arrow::flatbuf::Buffer>& flatbuf_buffers,
        int64_t& offset
    );

    [[nodiscard]] std::vector<org::apache::arrow::flatbuf::Buffer>
    get_buffers(const sparrow::record_batch& record_batch);

    void fill_compressed_buffers(
        const sparrow::arrow_proxy& arrow_proxy,
        std::vector<org::apache::arrow::flatbuf::Buffer>& flatbuf_compressed_buffers,
        int64_t& offset,
        const CompressionType compression_type,
        CompressionCache& cache
    );

    [[nodiscard]] std::vector<org::apache::arrow::flatbuf::Buffer> get_compressed_buffers(
        const sparrow::record_batch& record_batch,
        const CompressionType compression_type,
        CompressionCache& cache
    );

    [[nodiscard]] int64_t calculate_body_size(
        const sparrow::arrow_proxy& arrow_proxy,
        std::optional<CompressionType> compression = std::nullopt,
        std::optional<std::reference_wrapper<CompressionCache>> cache = std::nullopt
    );

    [[nodiscard]] int64_t calculate_body_size(
        const sparrow::record_batch& record_batch,
        std::optional<CompressionType> compression = std::nullopt,
        std::optional<std::reference_wrapper<CompressionCache>> cache = std::nullopt
    );

    [[nodiscard]] flatbuffers::FlatBufferBuilder get_record_batch_message_builder(
        const sparrow::record_batch& record_batch,
        std::optional<CompressionType> compression = std::nullopt,
        std::optional<std::reference_wrapper<CompressionCache>> cache = std::nullopt
    );

    [[nodiscard]] flatbuffers::FlatBufferBuilder get_dictionary_batch_message_builder(
        int64_t dictionary_id,
        const sparrow::record_batch& record_batch,
        bool is_delta = false,
        std::optional<CompressionType> compression = std::nullopt,
        std::optional<std::reference_wrapper<CompressionCache>> cache = std::nullopt
    );
 
    // Helper function to extract and parse the footer from Arrow IPC file data
    [[nodiscard]] SPARROW_IPC_API const org::apache::arrow::flatbuf::Footer* get_footer_from_file_data(std::span<const uint8_t> file_data);
}