[opt](point query) Reuse segments for point query row data lookup

be/src/service/point_query_executor.cpp

@@ -506,10 +506,10 @@ Status PointQueryExecutor::_lookup_row_key() {
         }
         // Get rowlocation and rowset, ctx._rowset_ptr will acquire wrap this ptr
         auto rowset_ptr = std::make_unique<RowsetSharedPtr>();
-        st = (_tablet->lookup_row_key(_row_read_ctxs[i]._primary_key, nullptr, false,
-                                      specified_rowsets, &location, INT32_MAX /*rethink?*/,
-                                      segment_caches, rowset_ptr.get(), false, nullptr,
-                                      &_profile_metrics.read_stats, nullptr, &io_ctx));
+        st = (_tablet->lookup_row_key(
+                _row_read_ctxs[i]._primary_key, nullptr, false, specified_rowsets, &location,
+                INT32_MAX /*rethink?*/, segment_caches, rowset_ptr.get(), false, nullptr,
+                &_profile_metrics.read_stats, nullptr, &io_ctx, &_row_read_ctxs[i]._segment));
         if (st.is<ErrorCode::KEY_NOT_FOUND>()) {
             continue;
         }
@@ -554,8 +554,8 @@ Status PointQueryExecutor::_lookup_row_data() {
                 io_ctx.remote_scan_cache_write_limiter = _remote_scan_cache_write_limiter.get();
                 RETURN_IF_ERROR(_tablet->lookup_row_data(
                         _row_read_ctxs[i]._primary_key, _row_read_ctxs[i]._row_location.value(),
-                        *(_row_read_ctxs[i]._rowset_ptr), _profile_metrics.read_stats, value,
-                        use_row_cache, &io_ctx));
+                        _row_read_ctxs[i]._segment, *(_row_read_ctxs[i]._rowset_ptr),
+                        _profile_metrics.read_stats, value, use_row_cache, &io_ctx));
                 // serialize value to block, currently only jsonb row format
                 RETURN_IF_ERROR(JsonbSerializeUtil::jsonb_to_columns(
                         _reusable->get_data_type_serdes(), value.data(), value.size(),

be/src/service/point_query_executor.h

@@ -335,6 +335,8 @@ class PointQueryExecutor {
         // rowset will be aquired during read
         // and released after used
         std::unique_ptr<RowsetSharedPtr, decltype(&release_rowset)> _rowset_ptr;
+        // Reuse the segment loaded during key lookup; release it before its acquired rowset.
+        segment_v2::SegmentSharedPtr _segment;
     };
 
     PTabletKeyLookupResponse* _response = nullptr;

be/src/storage/segment/segment_loader.cpp

@@ -21,6 +21,7 @@
 
 #include "common/config.h"
 #include "common/status.h"
+#include "cpp/sync_point.h"
 #include "storage/olap_define.h"
 #include "storage/rowset/beta_rowset.h"
 #include "storage/rowset/rowset_meta.h"
@@ -118,6 +119,7 @@ Status SegmentLoader::load_segments(const BetaRowsetSharedPtr& rowset,
                                     bool need_load_pk_index_and_bf,
                                     OlapReaderStatistics* index_load_stats,
                                     const io::IOContext* io_ctx) {
+    TEST_SYNC_POINT_CALLBACK("SegmentLoader::load_segments", rowset.get());
     if (cache_handle->is_inited()) {
         return Status::OK();
     }

be/src/storage/tablet/base_tablet.cpp

@@ -484,6 +484,26 @@ Status BaseTablet::lookup_row_data(const Slice& encoded_key, const RowLocation&
                                    RowsetSharedPtr input_rowset, OlapReaderStatistics& stats,
                                    std::string& values, bool write_to_cache,
                                    const io::IOContext* io_ctx) {
+    return _lookup_row_data(encoded_key, row_location, nullptr, std::move(input_rowset), stats,
+                            values, write_to_cache, io_ctx);
+}
+
+Status BaseTablet::lookup_row_data(const Slice& encoded_key, const RowLocation& row_location,
+                                   const segment_v2::SegmentSharedPtr& segment,
+                                   RowsetSharedPtr input_rowset, OlapReaderStatistics& stats,
+                                   std::string& values, bool write_to_cache,
+                                   const io::IOContext* io_ctx) {
+    DCHECK(segment != nullptr);
+    DCHECK_EQ(segment->id(), row_location.segment_id);
+    return _lookup_row_data(encoded_key, row_location, &segment, std::move(input_rowset), stats,
+                            values, write_to_cache, io_ctx);
+}
+
+Status BaseTablet::_lookup_row_data(const Slice& encoded_key, const RowLocation& row_location,
+                                    const segment_v2::SegmentSharedPtr* segment,
+                                    RowsetSharedPtr input_rowset, OlapReaderStatistics& stats,
+                                    std::string& values, bool write_to_cache,
+                                    const io::IOContext* io_ctx) const {
     MonotonicStopWatch watch;
     size_t row_size = 1;
     watch.start();
@@ -498,9 +518,14 @@ Status BaseTablet::lookup_row_data(const Slice& encoded_key, const RowLocation&
     SegmentCacheHandle segment_cache_handle;
     std::unique_ptr<segment_v2::ColumnIterator> column_iterator;
     const auto& column = *DORIS_TRY(tablet_schema->column(BeConsts::ROW_STORE_COL));
-    RETURN_IF_ERROR(_get_segment_column_iterator(rowset, row_location.segment_id, column,
-                                                 &segment_cache_handle, &column_iterator, stats,
-                                                 io_ctx));
+    if (segment != nullptr) {
+        RETURN_IF_ERROR(_init_segment_column_iterator(rowset, *segment, column, &column_iterator,
+                                                      stats, io_ctx));
+    } else {
+        RETURN_IF_ERROR(_get_segment_column_iterator(rowset, row_location.segment_id, column,
+                                                     &segment_cache_handle, &column_iterator, stats,
+                                                     io_ctx));
+    }
     // get and parse tuple row
     MutableColumnPtr column_ptr = ColumnString::create();
     std::vector<segment_v2::rowid_t> rowids {static_cast<segment_v2::rowid_t>(row_location.row_id)};
@@ -523,7 +548,8 @@ Status BaseTablet::lookup_row_key(const Slice& encoded_key, TabletSchema* latest
                                   std::vector<std::unique_ptr<SegmentCacheHandle>>& segment_caches,
                                   RowsetSharedPtr* rowset, bool with_rowid,
                                   std::string* encoded_seq_value, OlapReaderStatistics* stats,
-                                  DeleteBitmapPtr delete_bitmap, const io::IOContext* io_ctx) {
+                                  DeleteBitmapPtr delete_bitmap, const io::IOContext* io_ctx,
+                                  segment_v2::SegmentSharedPtr* segment) {
     SCOPED_BVAR_LATENCY(g_tablet_lookup_rowkey_latency);
     size_t seq_col_length = 0;
     // use the latest tablet schema to decide if the tablet has sequence column currently
@@ -608,6 +634,9 @@ Status BaseTablet::lookup_row_key(const Slice& encoded_key, TabletSchema* latest
             }
             TEST_SYNC_POINT_CALLBACK("BaseTablet::lookup_row_key:found", this, rs.get(),
                                      with_seq_col, s.code());
+            if (segment) {
+                *segment = segments[id];
+            }
             // find it and return
             return s;
         }

be/src/storage/tablet/base_tablet.h

@@ -182,6 +182,13 @@ class BaseTablet : public std::enable_shared_from_this<BaseTablet> {
     Status lookup_row_data(const Slice& encoded_key, const RowLocation& row_location,
                            RowsetSharedPtr rowset, OlapReaderStatistics& stats, std::string& values,
                            bool write_to_cache = false, const io::IOContext* io_ctx = nullptr);
+
+    // Reuse the segment returned by lookup_row_key. The caller must keep the segment and
+    // its acquired rowset alive until this method returns.
+    Status lookup_row_data(const Slice& encoded_key, const RowLocation& row_location,
+                           const segment_v2::SegmentSharedPtr& segment, RowsetSharedPtr rowset,
+                           OlapReaderStatistics& stats, std::string& values,
+                           bool write_to_cache = false, const io::IOContext* io_ctx = nullptr);
     // Lookup the row location of `encoded_key`, the function sets `row_location` on success.
     // NOTE: the method only works in unique key model with primary key index, you will got a
     //       not supported error in other data model.
@@ -193,7 +200,8 @@ class BaseTablet : public std::enable_shared_from_this<BaseTablet> {
                           std::string* encoded_seq_value = nullptr,
                           OlapReaderStatistics* stats = nullptr,
                           DeleteBitmapPtr tablet_delete_bitmap = nullptr,
-                          const io::IOContext* io_ctx = nullptr);
+                          const io::IOContext* io_ctx = nullptr,
+                          segment_v2::SegmentSharedPtr* segment = nullptr);
 
     // calc delete bitmap when flush memtable, use a fake version to calc
     // For example, cur max version is 5, and we use version 6 to calc but
@@ -371,6 +379,11 @@ class BaseTablet : public std::enable_shared_from_this<BaseTablet> {
                                                                const CaptureRowsetOps& options);
 
 protected:
+    Status _lookup_row_data(const Slice& encoded_key, const RowLocation& row_location,
+                            const segment_v2::SegmentSharedPtr* segment, RowsetSharedPtr rowset,
+                            OlapReaderStatistics& stats, std::string& values, bool write_to_cache,
+                            const io::IOContext* io_ctx) const;
+
     // Protected by the tablet header lock; discovery resumes after this version.
     Version _changelog_ttl_scan_cursor {-1, -1};