performance

native-sql-engine/cpp/src/codegen/common/sort_relation.h

@@ -71,6 +71,7 @@ class SortRelation {
     int64_t batch_remaining = (batch_length - 1) - offset_in_current_batch_;
     if (shift <= batch_remaining) {
       offset_in_current_batch_ = offset_in_current_batch_ + shift;
+      ClearShiftCache();
       return;
     }
     int64_t remaining = shift - batch_remaining;
@@ -85,18 +86,44 @@ class SortRelation {
           ArrayRelease(i);
         }
         offset_in_current_batch_ = remaining - 1;
+        ClearShiftCache();
         return;
       }
       remaining -= current_batch_length;
       batch_i++;
     }
   }
 
-  ArrayItemIndexS GetItemIndexWithShift(int shift) {
+  void ClearShiftCache() {
+    shift_cache_validity_.clear();
+    shift_cache_.clear();
+  }
+
+  void AddToShiftCache(int32_t shift, ArrayItemIndexS &index_s) {
+    GrowShiftCache(shift);
+    shift_cache_validity_[shift] = true;
+    shift_cache_[shift] = index_s;
+  }
+
+  void GrowShiftCache(int32_t shift) {
+    if (shift < shift_cache_validity_.size()) {
+      return;
+    }
+    while (shift >= shift_cache_validity_.size()) {
+      shift_cache_validity_.resize(static_cast<int64_t>(shift_cache_validity_.size() * 1.5), false);
+    }
+    shift_cache_.resize(shift_cache_validity_.size());
+  }
+
+  ArrayItemIndexS GetItemIndexWithShift(int32_t shift) {
+    if (shift < shift_cache_validity_.size() && shift_cache_validity_[shift]) {
+      return shift_cache_[shift];
+    }
     int64_t batch_length = lazy_in_->GetNumRowsOfBatch(requested_batches);
     int64_t batch_remaining = (batch_length - 1) - offset_in_current_batch_;
     if (shift <= batch_remaining) {
       ArrayItemIndexS s(requested_batches, offset_in_current_batch_ + shift);
+      AddToShiftCache(shift, s);
       return s;
     }
     int64_t remaining = shift - batch_remaining;
@@ -105,6 +132,7 @@ class SortRelation {
       int64_t current_batch_length = lazy_in_->GetNumRowsOfBatch(batch_i);
       if (remaining <= current_batch_length) {
         ArrayItemIndexS s(batch_i, remaining - 1);
+        AddToShiftCache(shift, s);
         return s;
       }
       remaining -= current_batch_length;
@@ -113,6 +141,9 @@ class SortRelation {
   }
 
   bool CheckRangeBound(int shift) {
+    if (shift < shift_cache_validity_.size()) {
+      return true;
+    }
     int64_t batch_length = lazy_in_->GetNumRowsOfBatch(requested_batches);
     if (batch_length == -1L) {
       return false;
@@ -190,6 +221,8 @@ class SortRelation {
   int64_t offset_in_current_batch_ = 0;
   int32_t requested_batches = 0;
   int range_cache_ = -1;
+  std::vector<ArrayItemIndexS> shift_cache_;
+  std::vector<bool> shift_cache_validity_;
   std::vector<std::shared_ptr<RelationColumn>> sort_relation_key_list_;
   std::vector<std::shared_ptr<RelationColumn>> sort_relation_payload_list_;
 };