Storage: Refactor DMFileReader

dbms/src/Storages/DeltaMerge/BitmapFilter/BitmapFilter.cpp

@@ -12,9 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <Common/Exception.h>
 #include <Storages/DeltaMerge/BitmapFilter/BitmapFilter.h>
 #include <Storages/DeltaMerge/DeltaMergeHelpers.h>
-#include <Storages/DeltaMerge/Segment.h>
 
 namespace DB::DM
 {
@@ -30,53 +30,46 @@ void BitmapFilter::set(BlockInputStreamPtr & stream)
     {
         FilterPtr f = nullptr;
         auto blk = stream->read(f, /*res_filter*/ true);
-        if (likely(blk))
-        {
-            set(blk.segmentRowIdCol(), f);
-        }
-        else
+        if (unlikely(!blk))
         {
             break;
         }
+
+        const auto & row_ids_col = blk.segmentRowIdCol();
+        const auto * v = toColumnVectorDataPtr<UInt32>(row_ids_col);
+        assert(v != nullptr); // the segmentRowIdCol must be a UInt32 column
+        set(std::span{v->data(), v->size()}, f);
     }
     stream->readSuffix();
 }
 
-void BitmapFilter::set(const ColumnPtr & col, const FilterPtr & f)
-{
-    const auto * v = toColumnVectorDataPtr<UInt32>(col);
-    set(v->data(), v->size(), f);
-}
-
-void BitmapFilter::set(const UInt32 * data, UInt32 size, const FilterPtr & f)
+void BitmapFilter::set(std::span<const UInt32> row_ids, const FilterPtr & f)
 {
-    if (size == 0)
+    if (row_ids.empty())
     {
         return;
     }
     if (!f)
     {
-        for (UInt32 i = 0; i < size; i++)
+        for (auto row_id : row_ids)
         {
-            UInt32 row_id = *(data + i);
             filter[row_id] = true;
         }
     }
     else
     {
-        RUNTIME_CHECK(size == f->size(), size, f->size());
-        for (UInt32 i = 0; i < size; i++)
+        RUNTIME_CHECK(row_ids.size() == f->size(), row_ids.size(), f->size());
+        for (UInt32 i = 0; i < row_ids.size(); i++)
         {
-            UInt32 row_id = *(data + i);
-            filter[row_id] = (*f)[i];
+            filter[row_ids[i]] = (*f)[i];
         }
     }
 }
 
-void BitmapFilter::set(UInt32 start, UInt32 limit)
+void BitmapFilter::set(UInt32 start, UInt32 limit, bool value)
 {
     RUNTIME_CHECK(start + limit <= filter.size(), start, limit, filter.size());
-    std::fill(filter.begin() + start, filter.begin() + start + limit, true);
+    std::fill_n(filter.begin() + start, limit, value);
 }
 
 bool BitmapFilter::get(IColumn::Filter & f, UInt32 start, UInt32 limit) const
@@ -127,4 +120,4 @@ size_t BitmapFilter::count() const
 {
     return std::count(filter.cbegin(), filter.cend(), true);
 }
-} // namespace DB::DM
+} // namespace DB::DM

dbms/src/Storages/DeltaMerge/BitmapFilter/BitmapFilter.h

@@ -17,6 +17,8 @@
 #include <Columns/IColumn.h>
 #include <DataStreams/IBlockInputStream.h>
 
+#include <span>
+
 namespace DB::DM
 {
 
@@ -25,13 +27,14 @@ class BitmapFilter
 public:
     BitmapFilter(UInt32 size_, bool default_value);
 
+    // Read blocks from `stream` and set the rows_id to be true according to the
+    // `segmentRowIdCol` in the block read from `stream`.
     void set(BlockInputStreamPtr & stream);
-    void set(const ColumnPtr & col, const FilterPtr & f);
-    void set(const UInt32 * data, UInt32 size, const FilterPtr & f);
-    void set(UInt32 start, UInt32 limit);
+    // f[start, satrt+limit) = value
+    void set(UInt32 start, UInt32 limit, bool value = true);
     // If return true, all data is match and do not fill the filter.
     bool get(IColumn::Filter & f, UInt32 start, UInt32 limit) const;
-    // filter[start, limit] & f -> f
+    // filter[start, satrt+limit) & f -> f
     void rangeAnd(IColumn::Filter & f, UInt32 start, UInt32 limit) const;
 
     void runOptimize();
@@ -40,9 +43,11 @@ class BitmapFilter
     size_t count() const;
 
 private:
+    void set(std::span<const UInt32> row_ids, const FilterPtr & f);
+
     std::vector<bool> filter;
     bool all_match;
 };
 
 using BitmapFilterPtr = std::shared_ptr<BitmapFilter>;
-} // namespace DB::DM
+} // namespace DB::DM

dbms/src/Storages/DeltaMerge/File/ColumnCache.cpp

@@ -18,15 +18,15 @@ namespace DB
 {
 namespace DM
 {
-RangeWithStrategys ColumnCache::getReadStrategy(size_t pack_id, size_t pack_count, ColId column_id)
+RangeWithStrategys ColumnCache::getReadStrategy(size_t start_pack_idx, size_t pack_count, ColId column_id)
 {
-    PackRange target_range{pack_id, pack_id + pack_count};
+    PackRange target_range{start_pack_idx, start_pack_idx + pack_count};
 
     RangeWithStrategys range_and_strategys;
-
-    Strategy strategy = Strategy::Unknown;
+    range_and_strategys.reserve(pack_count);
+    auto strategy = Strategy::Unknown;
     size_t range_start = 0;
-    for (size_t cursor = target_range.first; cursor < target_range.second; cursor++)
+    for (size_t cursor = target_range.first; cursor < target_range.second; ++cursor)
     {
         if (isPackInCache(cursor, column_id))
         {
@@ -36,7 +36,7 @@ RangeWithStrategys ColumnCache::getReadStrategy(size_t pack_id, size_t pack_coun
             }
             else if (strategy == Strategy::Disk)
             {
-                range_and_strategys.emplace_back(std::make_pair(PackRange{range_start, cursor}, Strategy::Disk));
+                range_and_strategys.emplace_back(PackRange{range_start, cursor}, Strategy::Disk);
             }
             range_start = cursor;
             strategy = Strategy::Memory;
@@ -45,7 +45,7 @@ RangeWithStrategys ColumnCache::getReadStrategy(size_t pack_id, size_t pack_coun
         {
             if (strategy == Strategy::Memory)
             {
-                range_and_strategys.emplace_back(std::make_pair(PackRange{range_start, cursor}, Strategy::Memory));
+                range_and_strategys.emplace_back(PackRange{range_start, cursor}, Strategy::Memory);
             }
             else if (strategy == Strategy::Disk)
             {
@@ -55,8 +55,53 @@ RangeWithStrategys ColumnCache::getReadStrategy(size_t pack_id, size_t pack_coun
             strategy = Strategy::Disk;
         }
     }
-    range_and_strategys.emplace_back(std::make_pair(PackRange{range_start, target_range.second}, strategy));
+    range_and_strategys.emplace_back(PackRange{range_start, target_range.second}, strategy);
+    range_and_strategys.shrink_to_fit();
+    return range_and_strategys;
+}
+
+RangeWithStrategys ColumnCache::getReadStrategy(
+    size_t start_pack_idx,
+    size_t pack_count,
+    const std::vector<size_t> & clean_read_pack_idx)
+{
+    PackRange target_range{start_pack_idx, start_pack_idx + pack_count};
 
+    RangeWithStrategys range_and_strategys;
+    range_and_strategys.reserve(pack_count);
+    auto strategy = Strategy::Unknown;
+    size_t range_start = 0;
+    for (size_t cursor = target_range.first; cursor < target_range.second; ++cursor)
+    {
+        if (std::find(clean_read_pack_idx.cbegin(), clean_read_pack_idx.cend(), cursor) != clean_read_pack_idx.cend())
+        {
+            if (strategy == Strategy::Memory)
+            {
+                continue;
+            }
+            else if (strategy == Strategy::Disk)
+            {
+                range_and_strategys.emplace_back(PackRange{range_start, cursor}, Strategy::Disk);
+            }
+            range_start = cursor;
+            strategy = Strategy::Memory;
+        }
+        else
+        {
+            if (strategy == Strategy::Memory)
+            {
+                range_and_strategys.emplace_back(PackRange{range_start, cursor}, Strategy::Memory);
+            }
+            else if (strategy == Strategy::Disk)
+            {
+                continue;
+            }
+            range_start = cursor;
+            strategy = Strategy::Disk;
+        }
+    }
+    range_and_strategys.emplace_back(PackRange{range_start, target_range.second}, strategy);
+    range_and_strategys.shrink_to_fit();
     return range_and_strategys;
 }
 

dbms/src/Storages/DeltaMerge/File/ColumnCache.h

@@ -46,7 +46,11 @@ class ColumnCache
 
     using RangeWithStrategy = std::pair<PackRange, ColumnCache::Strategy>;
     using RangeWithStrategys = std::vector<RangeWithStrategy>;
-    RangeWithStrategys getReadStrategy(size_t pack_id, size_t pack_count, ColId column_id);
+    RangeWithStrategys getReadStrategy(size_t start_pack_idx, size_t pack_count, ColId column_id);
+    static RangeWithStrategys getReadStrategy(
+        size_t start_pack_idx,
+        size_t pack_count,
+        const std::vector<size_t> & clean_read_pack_idx);
 
     void tryPutColumn(size_t pack_id, ColId column_id, const ColumnPtr & column, size_t rows_offset, size_t rows_count);