code: import logging
  import os
  import time
  from functools import wraps
  from pathlib import Path
  from random import random, seed
  
  import lancedb
  import pyarrow as pa
  import pyarrow.parquet as pq
  import typer
  from lancedb.db import LanceTable
  
  log_level = os.environ.get("LOG_LEVEL", "info")
  logging.basicConfig(
      level=getattr(logging, log_level.upper()),
      format="%(asctime)s %(levelname)s | %(processName)s %(name)s | %(message)s",
  )
  
  
  logger = logging.getLogger(__name__)
  app = typer.Typer()
  
  
  V_SIZE = 256
  
  DB_PATH = "benchmark"
  DB_TABLE = "vectors"
  DB_TABLE_SIZE = os.environ.get("DB_TABLE_SIZE", 100000)
  
  Q_PATH = "query"
  Q_SIZE = os.environ.get("Q_SIZE", 100)
  Q_V = "v.parquet"
  Q_KNN = "knn.parquet"
  Q_ANN = "ann.parquet"
  
  
  def timeit(func):
      @wraps(func)
      def f(*args, **kwargs):
          start_time = time.perf_counter()
          result = func(*args, **kwargs)
          end_time = time.perf_counter()
          total_time = end_time - start_time
          logger.info(f"{func.__name__} {args} done in {total_time:.2f} secs")
          return result
  
      return f
  
  
  def get_db():
      if int(os.environ["AZURE"]) == 0:
          f = Path(os.environ["DATA"])
          f.mkdir(parents=True, exist_ok=True)
          return lancedb.connect(f / DB_PATH)
      else:
          return lancedb.connect(
              f"az://{os.environ['AZURE_STORAGE_CONTAINER']}/{DB_PATH}"
          )
  
  
  def open_table(table: str):
      return LanceTable(get_db(), table)
  
  
  def get_q(what="v"):
      tables = {
          "v": Q_V,
          "knn": Q_KNN,
          "ann": Q_ANN,
      }
      f = Path(os.environ["DATA"]) / Q_PATH
      f.mkdir(parents=True, exist_ok=True)
      return f / tables[what]
  
  
  def gen_data(n: int, start=1):
      seed()
      for i in range(start, start + n):
          yield ({"id": i, "vector": list(random() for _ in range(V_SIZE))})
  
  
  @app.command()
  def db_init(n: int = DB_TABLE_SIZE):
      get_db().create_table(DB_TABLE, data=list(gen_data(n)))
  
  
  @app.command()
  def db_info():
      table = open_table(DB_TABLE)
      logger.debug(table.head(10))
  
  
  @app.command()
  def db_add(n: int, start: int):
      table = open_table(DB_TABLE)
      table.add(list(gen_data(n, start=start)))
  
  
  @app.command()
  def q_init(n: int = Q_SIZE):
      pq.write_table(pa.Table.from_pylist(list(gen_data(n))), get_q())
  
  
  @app.command()
  def q_info():
      logger.debug(pq.read_table(get_q()))
  
  
  @timeit
  def q_process(what: str):
      table = open_table(DB_TABLE)
      r = pa.Table.from_pylist(
          [
              {
                  "id": v["id"],
                  "neighbours": table.search(v["vector"])
                  .limit(10)
                  .select(["id"])
                  .to_arrow()["id"]
                  .to_pylist(),
              }
              for v in pq.read_table(get_q()).to_pylist()
          ]
      )
      pq.write_table(r, get_q(what))
  
  
  @app.command()
  @timeit
  def create_index():
      open_table(DB_TABLE).create_index(
          num_sub_vectors=8
      )  # TODO :avoid hard coded params
  
  
  @app.command()
  def q_knn():
      q_process("knn")
  
  
  @app.command()
  def q_ann():
      create_index()
      q_process("ann")
  
  
  if __name__ == "__main__":
      app()
  
apis: lancedb.connect
extract_api: [((248, 283), 'os.environ.get', 'os.environ.get', (['"""LOG_LEVEL"""', '"""info"""'], {}), "('LOG_LEVEL', 'info')\n", (262, 283), False, 'import os\n'), ((446, 473), 'logging.getLogger', 'logging.getLogger', (['__name__'], {}), '(__name__)\n', (463, 473), False, 'import logging\n'), ((480, 493), 'typer.Typer', 'typer.Typer', ([], {}), '()\n', (491, 493), False, 'import typer\n'), ((569, 608), 'os.environ.get', 'os.environ.get', (['"""DB_TABLE_SIZE"""', '(100000)'], {}), "('DB_TABLE_SIZE', 100000)\n", (583, 608), False, 'import os\n'), ((636, 665), 'os.environ.get', 'os.environ.get', (['"""Q_SIZE"""', '(100)'], {}), "('Q_SIZE', 100)\n", (650, 665), False, 'import os\n'), ((753, 764), 'functools.wraps', 'wraps', (['func'], {}), '(func)\n', (758, 764), False, 'from functools import wraps\n'), ((1693, 1699), 'random.seed', 'seed', ([], {}), '()\n', (1697, 1699), False, 'from random import random, seed\n'), ((814, 833), 'time.perf_counter', 'time.perf_counter', ([], {}), '()\n', (831, 833), False, 'import time\n'), ((892, 911), 'time.perf_counter', 'time.perf_counter', ([], {}), '()\n', (909, 911), False, 'import time\n'), ((1134, 1158), 'pathlib.Path', 'Path', (["os.environ['DATA']"], {}), "(os.environ['DATA'])\n", (1138, 1158), False, 'from pathlib import Path\n'), ((1219, 1247), 'lancedb.connect', 'lancedb.connect', (['(f / DB_PATH)'], {}), '(f / DB_PATH)\n', (1234, 1247), False, 'import lancedb\n'), ((1273, 1347), 'lancedb.connect', 'lancedb.connect', (['f"""az://{os.environ[\'AZURE_STORAGE_CONTAINER\']}/{DB_PATH}"""'], {}), '(f"az://{os.environ[\'AZURE_STORAGE_CONTAINER\']}/{DB_PATH}")\n', (1288, 1347), False, 'import lancedb\n'), ((1553, 1577), 'pathlib.Path', 'Path', (["os.environ['DATA']"], {}), "(os.environ['DATA'])\n", (1557, 1577), False, 'from pathlib import Path\n'), ((1778, 1786), 'random.random', 'random', ([], {}), '()\n', (1784, 1786), False, 'from random import random, seed\n')]