A YSQL version of the test case for https://github.com/yugabyte/yugabyte-db/issues/13885

ysql_limit_one_scan.py

# Dependencies:
# On CentOS you can install psycopg2 thus:
#
#  sudo yum install postgresql-libs
#  sudo yum install python-psycopg2

import psycopg2
from threading import Thread,Semaphore
import random
import time

# Test Params:
num_offsets=5000

cluster_ip="localhost"

connect_string="host={} dbname=yugabyte user=yugabyte port=5433".format(cluster_ip)

def create_table():
  conn = psycopg2.connect(connect_string)
  conn.set_session(autocommit=True)
  cur = conn.cursor()
  cur.execute("""DROP TABLE IF EXISTS msg_q""");
  print("Dropped (if exists): msg_q table")
  print("====================")
  cur.execute("""
        CREATE TABLE IF NOT EXISTS msg_q (
           id     text,
           from_offset integer,
           until_offset integer,
           PRIMARY KEY(id HASH, from_offset ASC)
        )
    """)
  print("Created msg_q table.")
  cur.close()
  conn.close()

def init_rows():
  conn = psycopg2.connect(connect_string)
  conn.set_session(autocommit=True)
  cur = conn.cursor();
  print("Inserting {} rows".format(num_offsets))

  for idx in range(num_offsets):
    cur.execute("""INSERT INTO msg_q(id, from_offset, until_offset) VALUES (%s, %s, %s)""",
                ("id-0", idx, idx+1))
  print("====================")
  cur.close()
  conn.close()


def delete_offsets_upto(start_inclusive, end_exclusive):
  conn = psycopg2.connect(connect_string)
  conn.set_session(autocommit=True)
  cur = conn.cursor();
  print("Deleting from_offsets [{} to {}).".format(start_inclusive, end_exclusive))

  for idx in range(start_inclusive, end_exclusive):
    cur.execute("""DELETE FROM msg_q WHERE id = %s and from_offset = %s""",
                ("id-0", idx))
  cur.close()
  conn.close()

def delete_all_but_two_rows():
  delete_offsets_upto(0, num_offsets - 2)

def delete_last_but_one_row():
  delete_offsets_upto(num_offsets - 2, num_offsets - 1)

def delete_last_row():
  delete_offsets_upto(num_offsets - 1, num_offsets)

def select_max_offset():

  conn = psycopg2.connect(connect_string)
  cur = conn.cursor()
  num_reads = 100
  s_delta = 0
  for idx in range(num_reads):
    t1 = time.time()
    cur.execute("SELECT from_offset, until_offset FROM msg_q WHERE id = %s ORDER BY from_offset DESC LIMIT 1", ("id-0", ))
    record = cur.fetchall()
    t2 = time.time()
    s_delta = s_delta + (t2 - t1)

  print("Avg SELECT Time: %s ms" % (((s_delta) * 1000) / num_reads))
  # Let's print one result record for debug/verification purposes
  print(record)
  print("====================")
  cur.close()
  conn.close()

# Main
create_table()
init_rows()
select_max_offset()
delete_all_but_two_rows()
select_max_offset()
delete_last_but_one_row()
select_max_offset()
delete_last_row()
select_max_offset()

$ python ~/notes/ysql_limit_one_scan.py
Dropped (if exists): msg_q table
====================
Created msg_q table.
Inserting 5000 rows
====================
Avg SELECT Time: 1.08699321747 ms
[(4999, 5000)]
====================
Deleting from_offsets [0 to 4998).
Avg SELECT Time: 1.22771978378 ms
[(4999, 5000)]
====================
Deleting from_offsets [4998 to 4999).
Avg SELECT Time: 26.1149573326 ms   ---> This should have been faster (~1ms) given that the row we want exists
[(4999, 5000)]
====================
Deleting from_offsets [4999 to 5000).
Avg SELECT Time: 26.5253329277 ms
[]
====================