stream_file_serializer.hpp

Go to the documentation of this file.

#pragma once
 
#include <cstddef>
#include <optional>
#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/dictionary_iteration.hpp"
#include "sparrow_ipc/dictionary_tracker.hpp"
#include "sparrow_ipc/magic_values.hpp"
#include "sparrow_ipc/serialize.hpp"
#include "sparrow_ipc/serialize_utils.hpp"
#include "sparrow_ipc/serializer_reserve.hpp"
 
namespace sparrow_ipc
{
    struct record_batch_block
    {
        int64_t offset;           
        int32_t metadata_length;  
        int64_t body_length;      
    };

    SPARROW_IPC_API size_t write_footer(
        const sparrow::record_batch& record_batch,
        const std::vector<record_batch_block>& dictionary_blocks,
        const std::vector<record_batch_block>& record_batch_blocks,
        any_output_stream& stream
    );

    class SPARROW_IPC_API stream_file_serializer
    {
    public:

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

        template <writable_stream TStream>
        stream_file_serializer(
            TStream& stream,
            const sparrow::record_batch& schema_batch,
            std::optional<CompressionType> compression = std::nullopt
        )
            : m_stream(stream)
            , m_compression(compression)
        {
            // Write file header magic
            m_stream.write(arrow_file_header_magic);
            m_stream.add_padding();
            m_header_written = true;
 
            // Establish schema
            m_schema_received = true;
            m_first_record_batch = schema_batch;
            m_dtypes = get_column_dtypes(schema_batch);
            serialize_schema_message(schema_batch, m_stream);
        }

        ~stream_file_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 write to a file serializer that has been ended");
            }
 
            // Write file header magic on first write
            if (!m_header_written)
            {
                m_stream.write(arrow_file_header_magic);
                m_stream.add_padding();
                m_header_written = true;
            }
 
            // 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_first_record_batch = *record_batches.begin();
                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 file serializer schema");
                }
 
                for_each_pending_dictionary(rb, m_dict_tracker, [&](const dictionary_info& dict_info)
                {
                    if (m_dict_tracker.is_emitted(dict_info.id) && !dict_info.is_delta)
                    {
                        throw std::runtime_error(
                            "Arrow file format does not support multiple non-delta dictionary batches "
                            "for the same dictionary id"
                        );
                    }
 
                    const int64_t dict_offset = static_cast<int64_t>(m_stream.size());
                    const auto dict_block_info = serialize_dictionary_batch(
                        dict_info.id,
                        dict_info.data,
                        dict_info.is_delta,
                        m_stream,
                        m_compression,
                        compressed_buffers_cache
                    );
                    m_dictionary_blocks.emplace_back(
                        dict_offset,
                        dict_block_info.metadata_length,
                        dict_block_info.body_length
                    );
                });
 
                // Offset is from the start of the file to the record batch message
                const int64_t offset = static_cast<int64_t>(m_stream.size());
 
                // Serialize and get block info
                const auto info = serialize_record_batch(rb, m_stream, m_compression, compressed_buffers_cache);
 
                m_record_batch_blocks.emplace_back(offset, info.metadata_length, info.body_length);
            }
        }

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

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

        void end();
 
        bool m_header_written{false};
        bool m_schema_received{false};
        std::optional<sparrow::record_batch> m_first_record_batch;
        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;
        std::vector<record_batch_block> m_dictionary_blocks;
        std::vector<record_batch_block> m_record_batch_blocks;
    };

    inline stream_file_serializer& end_file(stream_file_serializer& serializer)
    {
        serializer.end();
        return serializer;
    }
}