From 692970dc1748c091ce2e03cc80c2b53a069d3e16 Mon Sep 17 00:00:00 2001
From: Andrei Martsinchyk <amartsinchyk@yugabyte.com>
Date: Thu, 4 Apr 2024 09:22:16 -0700
Subject: [PATCH] [#21694] YSQL: allow arbitrary upper bound in backward scan
 requests

Summary:
Previously the function to find the target partition for backward scan
might work incorrectly if the upper bound wasn't a partition bound.
It seems like it always was the case when bounds were set for a backward scan
request. However, parallel backward scan uses the bounds and they may be
arbitrary keys.

Also, fix the test.

Correct target partition to start backward scan with valid upper bound
is the partition holding the bound, unless the bound equals to lower key
of the partition and upper bound is exclusive, when target partition is
the previous partition.
Jira: DB-10577

Test Plan: ./yb_build --cxx-test client_client-test --gtest_filter ClientTest.TestKeyRangeUpperBoundFiltering

Reviewers: arybochkin, timur

Reviewed By: arybochkin

Subscribers: ybase, yql

Tags: #jenkins-ready

Differential Revision: https://phorge.dev.yugabyte.com/D33590
---
 src/yb/client/client-test.cc | 36 ++++++++----------------------------
 src/yb/client/yb_op.cc       |  8 +++-----
 2 files changed, 11 insertions(+), 33 deletions(-)

diff --git a/src/yb/client/client-test.cc b/src/yb/client/client-test.cc
index b2c0dd0ee8b..9e3f37c2ce9 100644
--- a/src/yb/client/client-test.cc
+++ b/src/yb/client/client-test.cc
@@ -720,17 +720,18 @@ TEST_F(ClientTest, TestKeyRangeUpperBoundFiltering) {
   // General cases.
   for (bool is_inclusive : { true, false }) {
     for (size_t idx = 0; idx < partitions.size(); ++idx) {
-      // Special case, tested seprately at the end.
-      if (idx == 0 && !is_inclusive) {
-        continue;
-      }
-
       auto check_key = [&partitions, &req, idx, is_inclusive](const std::string& key) -> Status {
-        SCHECK_FORMAT((idx > 0 || is_inclusive), IllegalState, "", "");
+        SCHECK(!key.empty(), IllegalState, "Invalid key");
         req.clear_upper_bound();
         req.mutable_upper_bound()->set_key(key);
         req.mutable_upper_bound()->set_is_inclusive(is_inclusive);
-        auto* expected = is_inclusive ? &partitions[idx] : &partitions[idx - 1];
+        auto* expected = &partitions[idx];
+        // If the key is the same as the lower partition bound and exclusive, the previous partition
+        // is the correct target.
+        if (!is_inclusive && *expected == key) {
+          SCHECK_GT(idx, 0, IllegalState, "Invalid partitions, first partition key must be empty");
+          expected = &partitions[idx - 1];
+        }
         auto result = VERIFY_RESULT_REF(TEST_FindPartitionKeyByUpperBound(partitions, req));
         SCHECK_EQ(*expected, result, IllegalState, Format(
             "idx = $0, is_inclusive = $1, upper_bound = \"$2\"",
@@ -745,12 +746,6 @@ TEST_F(ClientTest, TestKeyRangeUpperBoundFiltering) {
           PartitionSchema::DecodeMultiColumnHashValue(partitions[idx + 1]) :
           std::numeric_limits<decltype(start)>::max()) - 1;
 
-      // 0. Special case: upper bound is empty means upper bound matches first partition start.
-      if (key.empty()) {
-        ASSERT_EQ(idx, 0);
-        ASSERT_OK(check_key(key));
-      }
-
       // 1. Upper bound matches partition start.
       ASSERT_OK(check_key(PartitionSchema::EncodeMultiColumnHashValue(start)));
 
@@ -762,21 +757,6 @@ TEST_F(ClientTest, TestKeyRangeUpperBoundFiltering) {
       ASSERT_OK(check_key(PartitionSchema::EncodeMultiColumnHashValue(middle)));
     }
   }
-
-  // Special case: upper_bound is exclusive and points to the first partition.
-  // Generates crash in DEBUG and returns InvalidArgument for release builds.
-  req.clear_upper_bound();
-  req.mutable_upper_bound()->set_key(partitions.front());
-  req.mutable_upper_bound()->set_is_inclusive(false);
-#ifndef NDEBUG
-  ASSERT_DEATH({
-#endif
-  auto result = TEST_FindPartitionKeyByUpperBound(partitions, req);
-  ASSERT_NOK(result);
-  ASSERT_TRUE(result.status().IsInvalidArgument());
-#ifndef NDEBUG
-  }, ".*Upper bound must not be exclusive when it points to the first partition.*");
-#endif
 }
 
 TEST_F(ClientTest, TestListTables) {
diff --git a/src/yb/client/yb_op.cc b/src/yb/client/yb_op.cc
index aa80838fa41..7927ed19032 100644
--- a/src/yb/client/yb_op.cc
+++ b/src/yb/client/yb_op.cc
@@ -125,11 +125,9 @@ Result<const PartitionKey&> FindPartitionKeyByUpperBound(
     return partitions.back();
   }
 
-  auto idx = FindPartitionStartIndex(
-      partitions, static_cast<std::string_view>(request.upper_bound().key()));
-  if (!request.upper_bound().is_inclusive()) {
-    RSTATUS_DCHECK_NE(idx, 0U, InvalidArgument,
-                      "Upper bound must not be exclusive when it points to the first partition.");
+  const auto& key = request.upper_bound().key();
+  auto idx = FindPartitionStartIndex(partitions, key);
+  if (!request.upper_bound().is_inclusive() && key == partitions[idx]) {
     --idx;
   }
   return partitions[idx];