[BugFix] Acquire rowsets at querying

be/src/exec/pipeline/scan/olap_scan_context.cpp

@@ -4,6 +4,7 @@
 
 #include "exec/vectorized/olap_scan_node.h"
 #include "exprs/vectorized/runtime_filter_bank.h"
+#include "storage/tablet.h"
 
 namespace starrocks::pipeline {
 
@@ -38,6 +39,31 @@ Status OlapScanContext::prepare(RuntimeState* state) {
 
 void OlapScanContext::close(RuntimeState* state) {
     _chunk_buffer.close();
+    for (const auto& rowsets_per_tablet : _tablet_rowsets) {
+        Rowset::release_readers(rowsets_per_tablet);
+    }
+}
+
+Status OlapScanContext::capture_tablet_rowsets(const std::vector<TInternalScanRange*>& olap_scan_ranges) {
+    _tablet_rowsets.resize(olap_scan_ranges.size());
+    _tablets.resize(olap_scan_ranges.size());
+    for (int i = 0; i < olap_scan_ranges.size(); ++i) {
+        auto* scan_range = olap_scan_ranges[i];
+
+        int64_t version = strtoul(scan_range->version.c_str(), nullptr, 10);
+        ASSIGN_OR_RETURN(TabletSharedPtr tablet, vectorized::OlapScanNode::get_tablet(scan_range));
+
+        // Capture row sets of this version tablet.
+        {
+            std::shared_lock l(tablet->get_header_lock());
+            RETURN_IF_ERROR(tablet->capture_consistent_rowsets(Version(0, version), &_tablet_rowsets[i]));
+            Rowset::acquire_readers(_tablet_rowsets[i]);
+        }
+
+        _tablets[i] = std::move(tablet);
+    }
+
+    return Status::OK();
 }
 
 Status OlapScanContext::parse_conjuncts(RuntimeState* state, const std::vector<ExprContext*>& runtime_in_filters,

be/src/exec/pipeline/scan/olap_scan_context.h

@@ -13,6 +13,10 @@
 namespace starrocks {
 
 class ScanNode;
+class Tablet;
+using TabletSharedPtr = std::shared_ptr<Tablet>;
+class Rowset;
+using RowsetSharedPtr = std::shared_ptr<Rowset>;
 
 namespace vectorized {
 class RuntimeFilterProbeCollector;
@@ -56,6 +60,10 @@ class OlapScanContext final : public ContextWithDependency {
     bool has_active_input() const;
     BalancedChunkBuffer& get_shared_buffer();
 
+    Status capture_tablet_rowsets(const std::vector<TInternalScanRange*>& olap_scan_ranges);
+    const std::vector<TabletSharedPtr>& tablets() const { return _tablets; }
+    const std::vector<std::vector<RowsetSharedPtr>>& tablet_rowsets() const { return _tablet_rowsets; };
+
 private:
     vectorized::OlapScanNode* _scan_node;
 
@@ -77,6 +85,13 @@ class OlapScanContext final : public ContextWithDependency {
     bool _shared_scan;                  // Enable shared_scan
 
     std::atomic<bool> _is_prepare_finished{false};
+
+    // The row sets of tablets will become stale and be deleted, if compaction occurs
+    // and these row sets aren't referenced, which will typically happen when the tablets
+    // of the left table are compacted at building the right hash table. Therefore, reference
+    // the row sets into _tablet_rowsets in the preparation phase to avoid the row sets being deleted.
+    std::vector<TabletSharedPtr> _tablets;
+    std::vector<std::vector<RowsetSharedPtr>> _tablet_rowsets;
 };
 
 // OlapScanContextFactory creates different contexts for each scan operator, if _shared_scan is false.

be/src/exec/pipeline/scan/olap_scan_prepare_operator.cpp

@@ -26,8 +26,10 @@ OlapScanPrepareOperator::~OlapScanPrepareOperator() {
 Status OlapScanPrepareOperator::prepare(RuntimeState* state) {
     RETURN_IF_ERROR(SourceOperator::prepare(state));
 
-    RETURN_IF_ERROR(_capture_tablet_rowsets());
-    return _ctx->prepare(state);
+    RETURN_IF_ERROR(_ctx->prepare(state));
+    RETURN_IF_ERROR(_ctx->capture_tablet_rowsets(_morsel_queue->olap_scan_ranges()));
+
+    return Status::OK();
 }
 
 void OlapScanPrepareOperator::close(RuntimeState* state) {
@@ -46,8 +48,8 @@ StatusOr<vectorized::ChunkPtr> OlapScanPrepareOperator::pull_chunk(RuntimeState*
     Status status = _ctx->parse_conjuncts(state, runtime_in_filters(), runtime_bloom_filters());
 
     _morsel_queue->set_key_ranges(_ctx->key_ranges());
-    _morsel_queue->set_tablets(_tablets);
-    _morsel_queue->set_tablet_rowsets(_tablet_rowsets);
+    _morsel_queue->set_tablets(_ctx->tablets());
+    _morsel_queue->set_tablet_rowsets(_ctx->tablet_rowsets());
 
     _ctx->set_prepare_finished();
     if (!status.ok()) {
@@ -58,28 +60,6 @@ StatusOr<vectorized::ChunkPtr> OlapScanPrepareOperator::pull_chunk(RuntimeState*
     return nullptr;
 }
 
-Status OlapScanPrepareOperator::_capture_tablet_rowsets() {
-    auto olap_scan_ranges = _morsel_queue->olap_scan_ranges();
-    _tablet_rowsets.resize(olap_scan_ranges.size());
-    _tablets.resize(olap_scan_ranges.size());
-    for (int i = 0; i < olap_scan_ranges.size(); ++i) {
-        auto* scan_range = olap_scan_ranges[i];
-
-        int64_t version = strtoul(scan_range->version.c_str(), nullptr, 10);
-        ASSIGN_OR_RETURN(TabletSharedPtr tablet, vectorized::OlapScanNode::get_tablet(scan_range));
-
-        // Capture row sets of this version tablet.
-        {
-            std::shared_lock l(tablet->get_header_lock());
-            RETURN_IF_ERROR(tablet->capture_consistent_rowsets(Version(0, version), &_tablet_rowsets[i]));
-        }
-
-        _tablets[i] = std::move(tablet);
-    }
-
-    return Status::OK();
-}
-
 /// OlapScanPrepareOperatorFactory
 OlapScanPrepareOperatorFactory::OlapScanPrepareOperatorFactory(int32_t id, int32_t plan_node_id,
                                                                vectorized::OlapScanNode* const scan_node,

be/src/exec/pipeline/scan/olap_scan_prepare_operator.h

@@ -29,16 +29,6 @@ class OlapScanPrepareOperator final : public SourceOperator {
     StatusOr<vectorized::ChunkPtr> pull_chunk(RuntimeState* state) override;
 
 private:
-    Status _capture_tablet_rowsets();
-
-private:
-    // The row sets of tablets will become stale and be deleted, if compaction occurs
-    // and these row sets aren't referenced, which will typically happen when the tablets
-    // of the left table are compacted at building the right hash table. Therefore, reference
-    // the row sets into _tablet_rowsets in the preparation phase to avoid the row sets being deleted.
-    std::vector<TabletSharedPtr> _tablets;
-    std::vector<std::vector<RowsetSharedPtr>> _tablet_rowsets;
-
     OlapScanContextPtr _ctx;
 };
 

be/src/exec/vectorized/olap_scan_node.cpp

@@ -204,6 +204,10 @@ Status OlapScanNode::close(RuntimeState* state) {
         release_large_columns<BinaryColumn>(runtime_state()->chunk_size() * 512);
     }
 
+    for (const auto& rowsets_per_tablet : _tablet_rowsets) {
+        Rowset::release_readers(rowsets_per_tablet);
+    }
+
     return ScanNode::close(state);
 }
 
@@ -680,6 +684,7 @@ Status OlapScanNode::_capture_tablet_rowsets() {
         {
             std::shared_lock l(tablet->get_header_lock());
             RETURN_IF_ERROR(tablet->capture_consistent_rowsets(Version(0, version), &_tablet_rowsets[i]));
+            Rowset::acquire_readers(_tablet_rowsets[i]);
         }
     }