serialize_utils.hpp

Go to the documentation of this file.

#pragma once
 
#include <ranges>
#include <vector>
 
#include <sparrow/record_batch.hpp>
 
#include "sparrow_ipc/any_output_stream.hpp"
#include "sparrow_ipc/compression.hpp"
#include "sparrow_ipc/config/config.hpp"
#include "sparrow_ipc/utils.hpp"
 
namespace sparrow_ipc
{
    SPARROW_IPC_API void
    serialize_schema_message(const sparrow::record_batch& record_batch, any_output_stream& stream);
    
    [[nodiscard]] SPARROW_IPC_API std::size_t
    calculate_schema_message_size(const sparrow::record_batch& record_batch);

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

    [[nodiscard]] SPARROW_IPC_API std::size_t
    calculate_dictionary_batch_message_size(
        int64_t dictionary_id,
        const sparrow::record_batch& record_batch,
        bool is_delta,
        std::optional<CompressionType> compression = std::nullopt,
        std::optional<std::reference_wrapper<CompressionCache>> cache = std::nullopt
    );

    template <std::ranges::input_range R>
        requires std::same_as<std::ranges::range_value_t<R>, sparrow::record_batch>
    [[nodiscard]] std::size_t calculate_total_serialized_size(const R& record_batches,
                                                              std::optional<CompressionType> compression = std::nullopt,
                                                              std::optional<std::reference_wrapper<CompressionCache>> cache = std::nullopt)
    {
        if (record_batches.empty())
        {
            return 0;
        }
 
        if (!utils::check_record_batches_consistency(record_batches))
        {
            throw std::invalid_argument("Record batches have inconsistent schemas");
        }
 
        // Calculate schema message size (only once)
        auto it = std::ranges::begin(record_batches);
        std::size_t total_size = calculate_schema_message_size(*it);
 
        // Calculate record batch message sizes
        for (const auto& record_batch : record_batches)
        {
            total_size += calculate_record_batch_message_size(record_batch, compression, cache);
        }
 
        return total_size;
    }

    SPARROW_IPC_API void fill_body(const sparrow::arrow_proxy& arrow_proxy, any_output_stream& stream,
                                   std::optional<CompressionType> compression = std::nullopt,
                                   std::optional<std::reference_wrapper<CompressionCache>> cache = std::nullopt);

    SPARROW_IPC_API void generate_body(const sparrow::record_batch& record_batch, any_output_stream& stream,
                                       std::optional<CompressionType> compression = std::nullopt,
                                       std::optional<std::reference_wrapper<CompressionCache>> cache = std::nullopt);
 
    SPARROW_IPC_API std::vector<sparrow::data_type> get_column_dtypes(const sparrow::record_batch& rb);
}