stream_file_serializer.hpp

Go to the documentation of this file.

#pragma once
 
#include <cstddef>
#include <numeric>
#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/magic_values.hpp"
#include "sparrow_ipc/serialize.hpp"
#include "sparrow_ipc/serialize_utils.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>& record_batch_blocks,
        any_output_stream& stream
    );
    
    [[nodiscard]] SPARROW_IPC_API std::vector<sparrow::record_batch>
    deserialize_file(std::span<const uint8_t> data);

    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)
        {
        }

        ~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 std::accumulate(
                           record_batches.begin(),
                           record_batches.end(),
                           m_stream.size(),
                           [&compressed_buffers_cache, this](size_t acc, const sparrow::record_batch& rb)
                           {
                               return acc + calculate_record_batch_message_size(rb, m_compression, compressed_buffers_cache);
                           }
                       )
                       + (m_schema_received ? 0 : calculate_schema_message_size(*record_batches.begin()));
            };
 
            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");
                }
                
                // 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;
        std::vector<record_batch_block> m_record_batch_blocks;
    };

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