ColumnCreator.cpp

/*
 *  Copyright (C) GridGain Systems. All Rights Reserved.
 *  _________        _____ __________________        _____
 *  __  ____/___________(_)______  /__  ____/______ ____(_)_______
 *  _  / __  __  ___/__  / _  __  / _  / __  _  __ `/__  / __  __ \
 *  / /_/ /  _  /    _  /  / /_/ /  / /_/ /  / /_/ / _  /  _  / / /
 *  \____/   /_/     /_/   \_,__/   \____/   \__,_/  /_/   /_/ /_/
 */

#include "ColumnCreator.h"
#include "DataTypes.h"
#include "granularity.h"
#include "compression.h"
#include <cassert>
#include <algorithm>

ColumnCreator::ColumnCreator(const std::string &baseFileName, ColumnInfo columnInfo)
    : columnInfo(columnInfo)
    , flushedElementsCount(0) {
    if (!isFixedSizeType(columnInfo.dataType))
        throw std::logic_error("Only fixed sizes are supported in column.");

    dataFile.exceptions(std::ios::failbit | std::ios::badbit);
    indexFile.exceptions(std::ios::failbit | std::ios::badbit);

    std::ios_base::openmode mode =
        std::ios_base::trunc | std::ios_base::in | std::ios_base::out | std::ios_base::binary;

    dataFile.open(baseFileName + ".bin", mode);
    indexFile.open(baseFileName + ".mrk", mode);
}

ColumnCreator::~ColumnCreator() {
    try {
        close();
    }
    catch (...) {
        /* Ignore exceptions in destructor. */
    }
}

void ColumnCreator::addValue(const std::optional<BytesView> &value) {
    assert(columnInfo.hasFixedSize());

    const size_t valueSize = columnInfo.getFixedSize();

    if (currentBlock.size() + valueSize > COLUMN_BLOCK_SIZE_HINT)
        flushCurrentBlock();

    if (value) {
        if (value->size() != valueSize)
            throw std::runtime_error("Size mismatch for a fixed-size column");

        currentBlock.insert(currentBlock.end(), value->begin(), value->end());

        if (columnInfo.nullable)
            nullMark.push_back(std::byte{0});
    }
    else {
        if (!columnInfo.nullable)
            throw std::runtime_error("Writing nullable data to not nullable column");

        currentBlock.insert(currentBlock.end(), valueSize, std::byte{0});

        nullMark.push_back(std::byte{1});
    }
}

void ColumnCreator::write(BytesView data, BytesView newNullMark) {
    assert(isFixedSizeType(columnInfo.dataType));

    if (!columnInfo.nullable) {
        auto it = std::find(std::begin(newNullMark), std::end(newNullMark), std::byte{1});
        if (it != std::end(newNullMark))
            throw std::runtime_error("Writing nullable data to not nullable column");
    }

    size_t ts = getTypeSize(columnInfo.dataType);

    size_t alignedBlockSize = (COLUMN_BLOCK_SIZE_HINT + ts - 1) / ts * ts;

    const std::byte *chunkLeft = data.data();

    size_t sizeLeft = data.size();

    size_t writtenElementsCount = 0;

    while (currentBlock.size() + sizeLeft > alignedBlockSize) {
        size_t chunkPartSize = alignedBlockSize - currentBlock.size();

        assert(chunkPartSize % ts == 0);

        currentBlock.insert(currentBlock.end(), chunkLeft, chunkLeft + chunkPartSize);

        chunkLeft += chunkPartSize;

        sizeLeft -= chunkPartSize;

        if (columnInfo.nullable) {
            auto elementsCount = chunkPartSize / ts;

            if (newNullMark.empty()) {
                nullMark.insert(std::end(nullMark), elementsCount, std::byte{0});
            }
            else {
                auto newNullMarkBlock = newNullMark.substr(writtenElementsCount, elementsCount);

                nullMark.insert(std::end(nullMark), std::begin(newNullMarkBlock), std::end(newNullMarkBlock));
            }

            writtenElementsCount += elementsCount;
        }

        flushCurrentBlock();
    }

    if (sizeLeft != 0) {
        assert(data.data() + data.size() == chunkLeft + sizeLeft);

        assert(sizeLeft % ts == 0);

        currentBlock.insert(currentBlock.end(), chunkLeft, chunkLeft + sizeLeft);

        if (columnInfo.nullable) {
            auto remainNullMarkBlock = newNullMark.substr(writtenElementsCount);

            nullMark.insert(std::end(nullMark), std::begin(remainNullMarkBlock), std::end(remainNullMarkBlock));
        }
    }
}

void ColumnCreator::flushCurrentBlock() {
    if (currentBlock.empty())
        return;

    size_t ts = getTypeSize(columnInfo.dataType);

    uint64_t granuleSize = ts * COLUMN_GRANULARITY;

    uint64_t idxItem[2] = {};

    idxItem[0] = dataFile.tellp(); // Offset of block in file.

    for (size_t granulatedOffset = (flushedElementsCount * ts) % granuleSize; granulatedOffset < currentBlock.size();
         granulatedOffset += ts * COLUMN_GRANULARITY) {
        idxItem[1] = granulatedOffset; // Offset of data inside uncompressed block.

        indexFile.write(reinterpret_cast<char *>(idxItem), sizeof(idxItem));
    }

    auto compressedBlock = compress(BytesView {currentBlock.data(), currentBlock.size()}, columnInfo.dataType);

    dataFile.write(reinterpret_cast<const char *>(&compressedBlock.header), sizeof(compressedBlock.header));

    StorageDataBlock compressedNullMarksBlock;

    // writing nullMarks header just after main data header for optimizations purpose
    if (columnInfo.nullable) {
        auto nullableBytesView = BytesView {nullMark.data(), nullMark.size()};

        compressedNullMarksBlock = compress(nullableBytesView, DATA_TYPE::UINT8);

        dataFile.write(
            reinterpret_cast<const char *>(&compressedNullMarksBlock.header), sizeof(compressedNullMarksBlock.header));
    }

    // writing main data
    dataFile.write(reinterpret_cast<const char *>(compressedBlock.data.get()),
        static_cast<std::streamsize>(compressedBlock.header.compressedSize));

    // writing nullMarks data
    if (columnInfo.nullable) {
        dataFile.write(reinterpret_cast<const char *>(compressedNullMarksBlock.data.get()),
            static_cast<std::streamsize>(compressedNullMarksBlock.header.compressedSize));
    }

    flushedElementsCount += currentBlock.size() / ts;

    currentBlock.clear();

    if (columnInfo.nullable)
        nullMark.clear();
}

void ColumnCreator::close() {
    flushCurrentBlock();

    if (dataFile.is_open())
        dataFile.close();

    if (indexFile.is_open())
        indexFile.close();
}