From ee4eb97c3fc3425e060853425d5008b937985a23 Mon Sep 17 00:00:00 2001
From: Anton Popov <pad11rus@gmail.com>
Date: Tue, 14 Apr 2020 04:27:27 +0300
Subject: [PATCH] in-memory parts: preparation

---
 .../MergeTree/IMergeTreeDataPartWriter.cpp    | 38 +++++++++
 .../MergeTree/IMergeTreeDataPartWriter.h      | 78 +++++++++++++++++++
 2 files changed, 116 insertions(+)
 create mode 100644 src/Storages/MergeTree/IMergeTreeDataPartWriter.cpp
 create mode 100644 src/Storages/MergeTree/IMergeTreeDataPartWriter.h

diff --git a/src/Storages/MergeTree/IMergeTreeDataPartWriter.cpp b/src/Storages/MergeTree/IMergeTreeDataPartWriter.cpp
new file mode 100644
index 00000000000..c5e20a8a95a
--- /dev/null
+++ b/src/Storages/MergeTree/IMergeTreeDataPartWriter.cpp
@@ -0,0 +1,38 @@
+#include <Storages/MergeTree/IMergeTreeDataPartWriter.h>
+
+namespace DB
+{
+
+IMergeTreeDataPartWriter::IMergeTreeDataPartWriter(
+    const MergeTreeData & storage_,
+    const MergeTreeWriterSettings & settings_)
+    : storage(storage_), settings(settings_) {}
+
+IMergeTreeDataPartWriter::IMergeTreeDataPartWriter(
+    const MergeTreeData & storage_,
+    const NamesAndTypesList & columns_list_,
+    const MergeTreeIndices & skip_indices_,
+    const MergeTreeIndexGranularity & index_granularity_,
+    const MergeTreeWriterSettings & settings_)
+    : storage(storage_)
+    , columns_list(columns_list_)
+    , skip_indices(skip_indices_)
+    , index_granularity(index_granularity_)
+    , settings(settings_) {}
+
+Columns IMergeTreeDataPartWriter::releaseIndexColumns()
+{
+    return Columns(
+        std::make_move_iterator(index_columns.begin()),
+        std::make_move_iterator(index_columns.end()));
+}
+
+void IMergeTreeDataPartWriter::next()
+{
+    current_mark = next_mark;
+    index_offset = next_index_offset;
+}
+
+IMergeTreeDataPartWriter::~IMergeTreeDataPartWriter() = default;
+
+}
diff --git a/src/Storages/MergeTree/IMergeTreeDataPartWriter.h b/src/Storages/MergeTree/IMergeTreeDataPartWriter.h
new file mode 100644
index 00000000000..1cab6ae4b60
--- /dev/null
+++ b/src/Storages/MergeTree/IMergeTreeDataPartWriter.h
@@ -0,0 +1,78 @@
+#pragma once
+
+#include <IO/WriteBufferFromFile.h>
+#include <IO/WriteBufferFromFileBase.h>
+#include <Compression/CompressedWriteBuffer.h>
+#include <IO/HashingWriteBuffer.h>
+#include <Storages/MergeTree/MergeTreeData.h>
+#include <DataStreams/IBlockOutputStream.h>
+#include <Storages/MergeTree/IMergeTreeDataPart.h>
+#include <Disks/IDisk.h>
+
+
+namespace DB
+{
+
+
+/// Writes data part to disk in different formats.
+/// Calculates and serializes primary and skip indices if needed.
+class IMergeTreeDataPartWriter : private boost::noncopyable
+{
+public:
+    IMergeTreeDataPartWriter(const MergeTreeData & storage_,
+        const MergeTreeWriterSettings & settings_);
+
+    IMergeTreeDataPartWriter(const MergeTreeData & storage_,
+        const NamesAndTypesList & columns_list_,
+        const MergeTreeIndices & skip_indices_,
+        const MergeTreeIndexGranularity & index_granularity_,
+        const MergeTreeWriterSettings & settings_);
+
+    virtual ~IMergeTreeDataPartWriter();
+
+    virtual void write(
+        const Block & block, const IColumn::Permutation * permutation = nullptr,
+        /* Blocks with already sorted index columns */
+        const Block & primary_key_block = {}, const Block & skip_indexes_block = {}) = 0;
+
+    virtual void calculateAndSerializePrimaryIndex(const Block & /* primary_index_block */) {}
+    virtual void calculateAndSerializeSkipIndices(const Block & /* skip_indexes_block */) {}
+
+    /// Shift mark and offset to prepare read next mark.
+    /// You must call it after calling write method and optionally
+    ///  calling calculations of primary and skip indices.
+    void next();
+
+    /// Count index_granularity for block and store in `index_granularity`
+    virtual void fillIndexGranularity(const Block & /* block */) {}
+
+    virtual void initSkipIndices() {}
+    virtual void initPrimaryIndex() {}
+
+    virtual void finishDataSerialization(IMergeTreeDataPart::Checksums & checksums, bool sync = false) = 0;
+    virtual void finishPrimaryIndexSerialization(MergeTreeData::DataPart::Checksums & /* checksums */) {}
+    virtual void finishSkipIndicesSerialization(MergeTreeData::DataPart::Checksums & /* checksums */) {}
+
+    Columns releaseIndexColumns();
+    const MergeTreeIndexGranularity & getIndexGranularity() const { return index_granularity; }
+    const MergeTreeIndices & getSkipIndices() { return skip_indices; }
+
+protected:
+    const MergeTreeData & storage;
+    NamesAndTypesList columns_list;
+    MergeTreeIndices skip_indices;
+    MergeTreeIndexGranularity index_granularity;
+    MergeTreeWriterSettings settings;
+
+    size_t current_mark = 0;
+
+    /// The offset to the first row of the block for which you want to write the index.
+    size_t index_offset = 0;
+
+    size_t next_mark = 0;
+    size_t next_index_offset = 0;
+
+    MutableColumns index_columns;
+};
+
+}