Merge pull request ClickHouse#63233 from CurtizJ/return-back-61551

Return back ClickHouse#61551 (More optimal loading of marks)

Author: CurtizJ

src/Storages/MergeTree/MarkRange.cpp

@@ -81,6 +81,11 @@ size_t MarkRanges::getNumberOfMarks() const
     return result;
 }
 
+bool MarkRanges::isOneRangeForWholePart(size_t num_marks_in_part) const
+{
+    return size() == 1 && front().begin == 0 && front().end == num_marks_in_part;
+}
+
 void MarkRanges::serialize(WriteBuffer & out) const
 {
     writeBinaryLittleEndian(this->size(), out);

src/Storages/MergeTree/MarkRange.h

@@ -36,6 +36,7 @@ struct MarkRanges : public std::deque<MarkRange>
     using std::deque<MarkRange>::deque; /// NOLINT(modernize-type-traits)
 
     size_t getNumberOfMarks() const;
+    bool isOneRangeForWholePart(size_t num_marks_in_part) const;
 
     void serialize(WriteBuffer & out) const;
     String describe() const;

src/Storages/MergeTree/MergeTreeIOSettings.h

@@ -44,6 +44,8 @@ struct MergeTreeReaderSettings
     bool enable_multiple_prewhere_read_steps = false;
     /// If true, try to lower size of read buffer according to granule size and compressed block size.
     bool adjust_read_buffer_size = true;
+    /// If true, it's allowed to read the whole part without reading marks.
+    bool can_read_part_without_marks = false;
 };
 
 struct MergeTreeWriterSettings

src/Storages/MergeTree/MergeTreeIndexReader.cpp

@@ -31,6 +31,8 @@ std::unique_ptr<MergeTreeReaderStream> makeIndexReader(
         load_marks_threadpool,
         /*num_columns_in_mark=*/ 1);
 
+    marks_loader->startAsyncLoad();
+
     return std::make_unique<MergeTreeReaderStreamSingleColumn>(
         part->getDataPartStoragePtr(),
         index->getFileName(), extension, marks_count,
@@ -65,6 +67,7 @@ MergeTreeIndexReader::MergeTreeIndexReader(
         mark_cache,
         uncompressed_cache,
         std::move(settings));
+
     version = index_format.version;
 
     stream->adjustRightMark(getLastMark(all_mark_ranges_));

src/Storages/MergeTree/MergeTreeMarksLoader.cpp

@@ -64,6 +64,10 @@ MergeTreeMarksLoader::MergeTreeMarksLoader(
     , read_settings(read_settings_)
     , num_columns_in_mark(num_columns_in_mark_)
     , load_marks_threadpool(load_marks_threadpool_)
+{
+}
+
+void MergeTreeMarksLoader::startAsyncLoad()
 {
     if (load_marks_threadpool)
         future = loadMarksAsync();
@@ -102,6 +106,8 @@ MergeTreeMarksGetterPtr MergeTreeMarksLoader::loadMarks()
 
 MarkCache::MappedPtr MergeTreeMarksLoader::loadMarksImpl()
 {
+    LOG_TEST(getLogger("MergeTreeMarksLoader"), "Loading marks from path {}", mrk_path);
+
     /// Memory for marks must not be accounted as memory usage for query, because they are stored in shared cache.
     MemoryTrackerBlockerInThread temporarily_disable_memory_tracker;
 
@@ -218,7 +224,9 @@ MarkCache::MappedPtr MergeTreeMarksLoader::loadMarksSync()
         }
     }
     else
+    {
         loaded_marks = loadMarksImpl();
+    }
 
     if (!loaded_marks)
     {

src/Storages/MergeTree/MergeTreeMarksLoader.h

@@ -50,6 +50,7 @@ class MergeTreeMarksLoader
 
     ~MergeTreeMarksLoader();
 
+    void startAsyncLoad();
     MergeTreeMarksGetterPtr loadMarks();
     size_t getNumColumns() const { return num_columns_in_mark; }
 

src/Storages/MergeTree/MergeTreeReaderCompact.cpp

@@ -48,6 +48,7 @@ MergeTreeReaderCompact::MergeTreeReaderCompact(
     , profile_callback(profile_callback_)
     , clock_type(clock_type_)
 {
+    marks_loader->startAsyncLoad();
 }
 
 void MergeTreeReaderCompact::fillColumnPositions()

src/Storages/MergeTree/MergeTreeReaderStream.cpp

@@ -13,6 +13,7 @@ namespace ErrorCodes
 {
     extern const int ARGUMENT_OUT_OF_BOUND;
     extern const int CANNOT_READ_ALL_DATA;
+    extern const int LOGICAL_ERROR;
 }
 
 MergeTreeReaderStream::MergeTreeReaderStream(
@@ -41,14 +42,17 @@ MergeTreeReaderStream::MergeTreeReaderStream(
 {
 }
 
+void MergeTreeReaderStream::loadMarks()
+{
+    if (!marks_getter)
+        marks_getter = marks_loader->loadMarks();
+}
+
 void MergeTreeReaderStream::init()
 {
     if (initialized)
         return;
 
-    initialized = true;
-    marks_getter = marks_loader->loadMarks();
-
     /// Compute the size of the buffer.
     auto [max_mark_range_bytes, sum_mark_range_bytes] = estimateMarkRangeBytes(all_mark_ranges);
 
@@ -110,11 +114,15 @@ void MergeTreeReaderStream::init()
         data_buffer = non_cached_buffer.get();
         compressed_data_buffer = non_cached_buffer.get();
     }
+
+    initialized = true;
 }
 
 void MergeTreeReaderStream::seekToMarkAndColumn(size_t row_index, size_t column_position)
 {
     init();
+    loadMarks();
+
     const auto & mark = marks_getter->getMark(row_index, column_position);
 
     try
@@ -193,7 +201,7 @@ CompressedReadBufferBase * MergeTreeReaderStream::getCompressedDataBuffer()
     return compressed_data_buffer;
 }
 
-size_t MergeTreeReaderStreamSingleColumn::getRightOffset(size_t right_mark) const
+size_t MergeTreeReaderStreamSingleColumn::getRightOffset(size_t right_mark)
 {
     /// NOTE: if we are reading the whole file, then right_mark == marks_count
     /// and we will use max_read_buffer_size for buffer size, thus avoiding the need to load marks.
@@ -202,7 +210,8 @@ size_t MergeTreeReaderStreamSingleColumn::getRightOffset(size_t right_mark) cons
     if (marks_count == 0)
         return 0;
 
-    assert(right_mark <= marks_count);
+    chassert(right_mark <= marks_count);
+    loadMarks();
 
     if (right_mark == 0)
         return marks_getter->getMark(right_mark, 0).offset_in_compressed_file;
@@ -281,9 +290,9 @@ size_t MergeTreeReaderStreamSingleColumn::getRightOffset(size_t right_mark) cons
     return file_size;
 }
 
-std::pair<size_t, size_t> MergeTreeReaderStreamSingleColumn::estimateMarkRangeBytes(const MarkRanges & mark_ranges) const
+std::pair<size_t, size_t> MergeTreeReaderStreamSingleColumn::estimateMarkRangeBytes(const MarkRanges & mark_ranges)
 {
-    assert(marks_getter != nullptr);
+    loadMarks();
 
     size_t max_range_bytes = 0;
     size_t sum_range_bytes = 0;
@@ -302,7 +311,34 @@ std::pair<size_t, size_t> MergeTreeReaderStreamSingleColumn::estimateMarkRangeBy
     return {max_range_bytes, sum_range_bytes};
 }
 
-size_t MergeTreeReaderStreamMultipleColumns::getRightOffsetOneColumn(size_t right_mark_non_included, size_t column_position) const
+size_t MergeTreeReaderStreamSingleColumnWholePart::getRightOffset(size_t right_mark)
+{
+    if (right_mark != marks_count)
+    {
+        throw Exception(ErrorCodes::LOGICAL_ERROR,
+            "Expected one right mark: {}, got: {}",
+            marks_count, right_mark);
+    }
+    return file_size;
+}
+
+std::pair<size_t, size_t> MergeTreeReaderStreamSingleColumnWholePart::estimateMarkRangeBytes(const MarkRanges & mark_ranges)
+{
+    if (!mark_ranges.isOneRangeForWholePart(marks_count))
+    {
+        throw Exception(ErrorCodes::LOGICAL_ERROR,
+            "Expected one mark range that covers the whole part, got: {}",
+            mark_ranges.describe());
+    }
+    return {file_size, file_size};
+}
+
+void MergeTreeReaderStreamSingleColumnWholePart::seekToMark(size_t)
+{
+    throw Exception(ErrorCodes::LOGICAL_ERROR, "MergeTreeReaderStreamSingleColumnWholePart cannot seek to marks");
+}
+
+size_t MergeTreeReaderStreamMultipleColumns::getRightOffsetOneColumn(size_t right_mark_non_included, size_t column_position)
 {
     /// NOTE: if we are reading the whole file, then right_mark == marks_count
     /// and we will use max_read_buffer_size for buffer size, thus avoiding the need to load marks.
@@ -311,7 +347,8 @@ size_t MergeTreeReaderStreamMultipleColumns::getRightOffsetOneColumn(size_t righ
     if (marks_count == 0)
         return 0;
 
-    assert(right_mark_non_included <= marks_count);
+    chassert(right_mark_non_included <= marks_count);
+    loadMarks();
 
     if (right_mark_non_included == 0)
         return marks_getter->getMark(right_mark_non_included, column_position).offset_in_compressed_file;
@@ -347,9 +384,9 @@ size_t MergeTreeReaderStreamMultipleColumns::getRightOffsetOneColumn(size_t righ
 }
 
 std::pair<size_t, size_t>
-MergeTreeReaderStreamMultipleColumns::estimateMarkRangeBytesOneColumn(const MarkRanges & mark_ranges, size_t column_position) const
+MergeTreeReaderStreamMultipleColumns::estimateMarkRangeBytesOneColumn(const MarkRanges & mark_ranges, size_t column_position)
 {
-    assert(marks_getter != nullptr);
+    loadMarks();
 
     /// As a maximal range we return the maximal size of a whole stripe.
     size_t max_range_bytes = 0;
@@ -386,8 +423,9 @@ MergeTreeReaderStreamMultipleColumns::estimateMarkRangeBytesOneColumn(const Mark
     return {max_range_bytes, sum_range_bytes};
 }
 
-MarkInCompressedFile MergeTreeReaderStreamMultipleColumns::getStartOfNextStripeMark(size_t row_index, size_t column_position) const
+MarkInCompressedFile MergeTreeReaderStreamMultipleColumns::getStartOfNextStripeMark(size_t row_index, size_t column_position)
 {
+    loadMarks();
     const auto & current_mark = marks_getter->getMark(row_index, column_position);
 
     if (marks_getter->getNumColumns() == 1)
@@ -434,27 +472,27 @@ MarkInCompressedFile MergeTreeReaderStreamMultipleColumns::getStartOfNextStripeM
     return marks_getter->getMark(mark_index + 1, column_position + 1);
 }
 
-size_t MergeTreeReaderStreamOneOfMultipleColumns::getRightOffset(size_t right_mark_non_included) const
+size_t MergeTreeReaderStreamOneOfMultipleColumns::getRightOffset(size_t right_mark_non_included)
 {
     return getRightOffsetOneColumn(right_mark_non_included, column_position);
 }
 
-std::pair<size_t, size_t> MergeTreeReaderStreamOneOfMultipleColumns::estimateMarkRangeBytes(const MarkRanges & mark_ranges) const
+std::pair<size_t, size_t> MergeTreeReaderStreamOneOfMultipleColumns::estimateMarkRangeBytes(const MarkRanges & mark_ranges)
 {
     return estimateMarkRangeBytesOneColumn(mark_ranges, column_position);
 }
 
-size_t MergeTreeReaderStreamAllOfMultipleColumns::getRightOffset(size_t right_mark_non_included) const
+size_t MergeTreeReaderStreamAllOfMultipleColumns::getRightOffset(size_t right_mark_non_included)
 {
     return getRightOffsetOneColumn(right_mark_non_included, marks_loader->getNumColumns() - 1);
 }
 
-std::pair<size_t, size_t> MergeTreeReaderStreamAllOfMultipleColumns::estimateMarkRangeBytes(const MarkRanges & mark_ranges) const
+std::pair<size_t, size_t> MergeTreeReaderStreamAllOfMultipleColumns::estimateMarkRangeBytes(const MarkRanges & mark_ranges)
 {
     size_t max_range_bytes = 0;
     size_t sum_range_bytes = 0;
 
-    for (size_t i = 0; i < marks_getter->getNumColumns(); ++i)
+    for (size_t i = 0; i < marks_loader->getNumColumns(); ++i)
     {
         auto [current_max, current_sum] = estimateMarkRangeBytesOneColumn(mark_ranges, i);
 

src/Storages/MergeTree/MergeTreeReaderStream.h

@@ -40,6 +40,7 @@ class MergeTreeReaderStream
     /// Seeks to exact mark in file.
     void seekToMarkAndColumn(size_t row_index, size_t column_position);
 
+    /// Seeks to the start of the file.
     void seekToStart();
 
     /**
@@ -53,11 +54,11 @@ class MergeTreeReaderStream
 
 private:
     /// Returns offset in file up to which it's needed to read file to read all rows up to @right_mark mark.
-    virtual size_t getRightOffset(size_t right_mark) const = 0;
+    virtual size_t getRightOffset(size_t right_mark) = 0;
 
     /// Returns estimated max amount of bytes to read among mark ranges (which is used as size for read buffer)
     /// and total amount of bytes to read in all mark ranges.
-    virtual std::pair<size_t, size_t> estimateMarkRangeBytes(const MarkRanges & mark_ranges) const = 0;
+    virtual std::pair<size_t, size_t> estimateMarkRangeBytes(const MarkRanges & mark_ranges) = 0;
 
     const ReadBufferFromFileBase::ProfileCallback profile_callback;
     const clockid_t clock_type;
@@ -80,6 +81,7 @@ class MergeTreeReaderStream
 
 protected:
     void init();
+    void loadMarks();
 
     const MergeTreeReaderSettings settings;
     const size_t marks_count;
@@ -100,11 +102,25 @@ class MergeTreeReaderStreamSingleColumn : public MergeTreeReaderStream
     {
     }
 
-    size_t getRightOffset(size_t right_mark_non_included) const override;
-    std::pair<size_t, size_t> estimateMarkRangeBytes(const MarkRanges & mark_ranges) const override;
+    size_t getRightOffset(size_t right_mark_non_included) override;
+    std::pair<size_t, size_t> estimateMarkRangeBytes(const MarkRanges & mark_ranges) override;
     void seekToMark(size_t row_index) override { seekToMarkAndColumn(row_index, 0); }
 };
 
+class MergeTreeReaderStreamSingleColumnWholePart : public MergeTreeReaderStream
+{
+public:
+    template <typename... Args>
+    explicit MergeTreeReaderStreamSingleColumnWholePart(Args &&... args)
+        : MergeTreeReaderStream{std::forward<Args>(args)...}
+    {
+    }
+
+    size_t getRightOffset(size_t right_mark_non_included) override;
+    std::pair<size_t, size_t> estimateMarkRangeBytes(const MarkRanges & mark_ranges) override;
+    void seekToMark(size_t row_index) override;
+};
+
 /// Base class for reading from file that contains multiple columns.
 /// It is used to read from compact parts.
 /// See more details about data layout in MergeTreeDataPartCompact.h.
@@ -118,9 +134,9 @@ class MergeTreeReaderStreamMultipleColumns : public MergeTreeReaderStream
     }
 
 protected:
-    size_t getRightOffsetOneColumn(size_t right_mark_non_included, size_t column_position) const;
-    std::pair<size_t, size_t> estimateMarkRangeBytesOneColumn(const MarkRanges & mark_ranges, size_t column_position) const;
-    MarkInCompressedFile getStartOfNextStripeMark(size_t row_index, size_t column_position) const;
+    size_t getRightOffsetOneColumn(size_t right_mark_non_included, size_t column_position);
+    std::pair<size_t, size_t> estimateMarkRangeBytesOneColumn(const MarkRanges & mark_ranges, size_t column_position);
+    MarkInCompressedFile getStartOfNextStripeMark(size_t row_index, size_t column_position);
 };
 
 /// Class for reading a single column from file that contains multiple columns
@@ -135,8 +151,8 @@ class MergeTreeReaderStreamOneOfMultipleColumns : public MergeTreeReaderStreamMu
     {
     }
 
-    size_t getRightOffset(size_t right_mark_non_included) const override;
-    std::pair<size_t, size_t> estimateMarkRangeBytes(const MarkRanges & mark_ranges) const override;
+    size_t getRightOffset(size_t right_mark_non_included) override;
+    std::pair<size_t, size_t> estimateMarkRangeBytes(const MarkRanges & mark_ranges) override;
     void seekToMark(size_t row_index) override { seekToMarkAndColumn(row_index, column_position); }
 
 private:
@@ -154,8 +170,8 @@ class MergeTreeReaderStreamAllOfMultipleColumns : public MergeTreeReaderStreamMu
     {
     }
 
-    size_t getRightOffset(size_t right_mark_non_included) const override;
-    std::pair<size_t, size_t> estimateMarkRangeBytes(const MarkRanges & mark_ranges) const override;
+    size_t getRightOffset(size_t right_mark_non_included) override;
+    std::pair<size_t, size_t> estimateMarkRangeBytes(const MarkRanges & mark_ranges) override;
     void seekToMark(size_t row_index) override { seekToMarkAndColumn(row_index, 0); }
 };