ClickHouse/src/Columns/ColumnStringHelpers.h

#pragma once

#include <cstring>
#include <cassert>

#include <Columns/IColumn.h>
#include <Common/PODArray.h>
#include <Columns/ColumnString.h>
#include <Columns/ColumnFixedString.h>
#include <IO/WriteBufferFromVector.h>

namespace DB
{

namespace ErrorCodes
{
extern const int TOO_LARGE_STRING_SIZE;
}

namespace ColumnStringHelpers
{

/** Simplifies writing data to the ColumnString or ColumnFixedString via WriteBuffer.
 *
 *  Take care of little subtle details, like padding or proper offsets.
 */
template <typename ColumnType>
class WriteHelper
{
    ColumnType & col;
    WriteBufferFromVector<typename ColumnType::Chars> buffer;
    size_t prev_row_buffer_size = 0;

    static ColumnType & resizeColumn(ColumnType & column, size_t rows)
    {
        if constexpr (std::is_same_v<ColumnType, ColumnFixedString>)
            column.resize(rows);
        else
        {
            column.getOffsets().reserve(rows);
            /// Using coefficient 2 for initial size is arbitrary.
            column.getChars().resize(rows * 2);
        }
        return column;
    }

public:
    WriteHelper(ColumnType & col_, size_t expected_rows)
        : col(resizeColumn(col_, expected_rows))
        , buffer(col.getChars())
    {}

    ~WriteHelper() = default;

    void finalize()
    {
        buffer.finalize();
    }

    auto & getWriteBuffer()
    {
        return buffer;
    }

    inline void rowWritten()
    {
        if constexpr (std::is_same_v<ColumnType, ColumnFixedString>)
        {
            if (buffer.count() > prev_row_buffer_size + col.getN())
                throw Exception(
                        ErrorCodes::TOO_LARGE_STRING_SIZE,
                        "Too large string for FixedString column");

            // Pad with zeroes on the right to maintain FixedString invariant.
            const auto excess_bytes = buffer.count() % col.getN();
            const auto fill_bytes = col.getN() - excess_bytes;
            writeChar(0, fill_bytes, buffer);
        }
        else
        {
            writeChar(0, buffer);
            col.getOffsets().push_back(buffer.count());
        }

        prev_row_buffer_size = buffer.count();
    }
};

}

}