# SPDX-License-Identifier: Apache-2.0

# SPDX-FileCopyrightText: Copyright The Lance Authors

from datetime import datetime

import lance

import pyarrow.parquet as pq

# This file compares the performance of lance v1 and v2 on the lineitem dataset,

# specifically for random access scans

tab = pq.read_table("~/lineitemsf1.snappy.parquet")

dsv1 = lance.write_dataset(tab, "/tmp/lineitem.lancev1", data_storage_version="2.0")

dsv2 = lance.write_dataset(tab, "/tmp/lineitem.lancev2", data_storage_version="2.1")

dsv1 = lance.dataset("/tmp/lineitem.lancev1")

dsv2 = lance.dataset("/tmp/lineitem.lancev2")

start = datetime.now()

dsv1.to_table(filter="l_shipmode = 'FOB'", limit=10000)

duration = (datetime.now() - start).total_seconds()

print(f"V1 query time: {duration}s")

start = datetime.now()

dsv2.to_table(filter="l_shipmode = 'FOB'", limit=10000)

duration = (datetime.now() - start).total_seconds()

print(f"V2 query time: {duration}s")

start = datetime.now()

dsv1.take([1, 40, 100, 130, 200])

duration = (datetime.now() - start).total_seconds()

print(f"V1 query time: {duration}s")

start = datetime.now()

dsv2.take([1, 40, 100, 130, 200])

duration = (datetime.now() - start).total_seconds()

print(f"V2 query time: {duration}s")