serializer.hpp

Go to the documentation of this file.

#pragma once
 
#include <sparrow/record_batch.hpp>
 
#include "sparrow_ipc/any_output_stream.hpp"
#include "sparrow_ipc/compression.hpp"
#include "sparrow_ipc/dictionary_iteration.hpp"
#include "sparrow_ipc/dictionary_tracker.hpp"
#include "sparrow_ipc/serialize.hpp"
#include "sparrow_ipc/serialize_utils.hpp"
#include "sparrow_ipc/serializer_reserve.hpp"
 
namespace sparrow_ipc
{
    class SPARROW_IPC_API serializer
    {
    public:

        template <writable_stream TStream>
        serializer(TStream& stream, std::optional<CompressionType> compression = std::nullopt)
            : m_stream(stream)
            , m_compression(compression)
        {
        }

        template <writable_stream TStream>
        serializer(
            TStream& stream,
            const sparrow::record_batch& schema_batch,
            std::optional<CompressionType> compression = std::nullopt
        )
            : m_stream(stream)
            , m_compression(compression)
            , m_schema_received(true)
            , m_dtypes(get_column_dtypes(schema_batch))
        {
            serialize_schema_message(schema_batch, m_stream);
        }

        ~serializer();

        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)
        {
            CompressionCache compressed_buffers_cache;
            if (std::ranges::empty(record_batches))
            {
                return;
            }
 
            if (m_ended)
            {
                throw std::runtime_error("Cannot append to a serializer that has been ended");
            }
 
            // NOTE `reserve_function` is making us store a cache for the compressed buffers at this level.
            // The benefit of capacity allocation should be evaluated vs storing a cache of compressed buffers
            // of record batches.
            const auto reserve_function = [&record_batches, &compressed_buffers_cache, this]()
            {
                return calculate_serializer_reserve_size(
                    record_batches,
                    m_stream.size(),
                    m_schema_received,
                    m_compression,
                    m_dict_tracker,
                    compressed_buffers_cache
                );
            };
 
            m_stream.reserve(reserve_function);
 
            if (!m_schema_received)
            {
                m_schema_received = true;
                m_dtypes = get_column_dtypes(*record_batches.begin());
                serialize_schema_message(*record_batches.begin(), m_stream);
            }
 
            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");
                }
 
                for_each_pending_dictionary(rb, m_dict_tracker, [&](const dictionary_info& dict_info)
                {
                    serialize_dictionary_batch(
                        dict_info.id,
                        dict_info.data,
                        dict_info.is_delta,
                        m_stream,
                        m_compression,
                        compressed_buffers_cache
                    );
                });
 
                serialize_record_batch(rb, m_stream, m_compression, compressed_buffers_cache);
            }
        }

        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>
        serializer& operator<<(const R& record_batches)
        {
            write(record_batches);
            return *this;
        }

        serializer& operator<<(serializer& (*manip)(serializer&) )
        {
            return manip(*this);
        }

        void end();
 
    private:
 
        static std::vector<sparrow::data_type> get_column_dtypes(const sparrow::record_batch& rb);
 
        bool m_schema_received{false};
        std::vector<sparrow::data_type> m_dtypes;
        any_output_stream m_stream;
        bool m_ended{false};
        std::optional<CompressionType> m_compression;
        dictionary_tracker m_dict_tracker;
    };
 
    inline serializer& end_stream(serializer& serializer)
    {
        serializer.end();
        return serializer;
    }
}