chunk_memory_serializer.hpp

Go to the documentation of this file.

#pragma once
 
#include <numeric>
#include <optional>
#include <ranges>
#include <stdexcept>
#include <vector>
 
#include <sparrow/record_batch.hpp>
 
#include "sparrow_ipc/any_output_stream.hpp"
#include "sparrow_ipc/chunk_memory_output_stream.hpp"
#include "sparrow_ipc/compression.hpp"
#include "sparrow_ipc/config/config.hpp"
#include "sparrow_ipc/memory_output_stream.hpp"
#include "sparrow_ipc/serialize.hpp"
#include "sparrow_ipc/serialize_utils.hpp"
 
namespace sparrow_ipc
{
    class SPARROW_IPC_API chunk_serializer
    {
    public:

        chunk_serializer(chunked_memory_output_stream<std::vector<std::vector<uint8_t>>>& stream, std::optional<CompressionType> compression = std::nullopt);

        void write(const sparrow::record_batch& rb);

        template <std::ranges::input_range R>
            requires std::same_as<std::ranges::range_value_t<R>, sparrow::record_batch>
        void write(const R& record_batches);

        chunk_serializer& operator<<(const sparrow::record_batch& rb);

        template <std::ranges::input_range R>
            requires std::same_as<std::ranges::range_value_t<R>, sparrow::record_batch>
        chunk_serializer& operator<<(const R& record_batches);

        void end();
 
    private:
 
        bool m_schema_received{false};
        std::vector<sparrow::data_type> m_dtypes;
        chunked_memory_output_stream<std::vector<std::vector<uint8_t>>>* m_pstream;
        bool m_ended{false};
        std::optional<CompressionType> m_compression;
    };
 
    // Implementation
 
    template <std::ranges::input_range R>
        requires std::same_as<std::ranges::range_value_t<R>, sparrow::record_batch>
    void chunk_serializer::write(const R& record_batches)
    {
        if (m_ended)
        {
            throw std::runtime_error("Cannot append record batches to a serializer that has been ended");
        }
 
        m_pstream->reserve((m_schema_received ? 0 : 1) + m_pstream->size() + record_batches.size());
 
        if (!m_schema_received)
        {
            m_schema_received = true;
            m_dtypes = get_column_dtypes(*record_batches.begin());
            std::vector<uint8_t> schema_buffer;
            memory_output_stream stream(schema_buffer);
            any_output_stream astream(stream);
            serialize_schema_message(*record_batches.begin(), astream);
            m_pstream->write(std::move(schema_buffer));
        }
 
        for (const auto& rb : record_batches)
        {
            if (get_column_dtypes(rb) != m_dtypes)
            {
                throw std::invalid_argument("Record batch schema does not match serializer schema");
            }
            std::vector<uint8_t> buffer;
            memory_output_stream stream(buffer);
            any_output_stream astream(stream);
            CompressionCache compressed_buffers_cache;
            serialize_record_batch(rb, astream, m_compression, compressed_buffers_cache);
            m_pstream->write(std::move(buffer));
        }
    }
 
    inline chunk_serializer& chunk_serializer::operator<<(const sparrow::record_batch& rb)
    {
        write(rb);
        return *this;
    }
 
    template <std::ranges::input_range R>
        requires std::same_as<std::ranges::range_value_t<R>, sparrow::record_batch>
    chunk_serializer& chunk_serializer::operator<<(const R& record_batches)
    {
        write(record_batches);
        return *this;
    }
}