# SPDX-License-Identifier: Apache-2.0

# SPDX-FileCopyrightText: Copyright The Lance Authors

import shutil

from pathlib import Path

import lance

import pyarrow as pa

import pytest

from lance._datagen import rand_batches

N_DIMS = 768

KiB = 1024

MiB = 1024 * KiB

GiB = 1024 * MiB

# Mostly vector data. One id column, a caption, and an embedding vector

def create_captioned_image_data(num_bytes):

schema = pa.schema(

[

pa.field("int32", pa.int32()),

pa.field("text", pa.utf8()),

pa.field("vector", pa.list_(pa.float32(), N_DIMS)),

]

)

return schema, rand_batches(

schema, num_batches=8, batch_size_bytes=int(num_bytes / 8)

)

# Purely scalar data (schema based on TPC-H lineitem table)

def create_scalar_data(num_bytes):

schema = pa.schema(

[

pa.field("l_orderkey", pa.int64()),

pa.field("l_partkey", pa.int64()),

pa.field("l_suppkey", pa.int64()),

pa.field("l_linenumber", pa.int64()),

pa.field("l_quantity", pa.float64()),

pa.field("l_extendedprice", pa.float64()),

pa.field("l_discount", pa.float64()),

pa.field("l_tax", pa.float64()),

pa.field("l_returnflag", pa.utf8()),

pa.field("l_linestatus", pa.utf8()),

pa.field("l_shipdate", pa.date32()),

pa.field("l_commitdate", pa.date32()),

pa.field("l_receiptdate", pa.date32()),

pa.field("l_shipinstruct", pa.utf8()),

pa.field("l_shipmode", pa.utf8()),

pa.field("l_comment", pa.utf8()),

]

)

return schema, rand_batches(

schema, num_batches=8, batch_size_bytes=int(num_bytes / 8)

)

def do_write_dataset(data, tmp_path, schema):

shutil.rmtree(tmp_path)

lance.write_dataset(data, tmp_path, schema=schema)

def write_dataset_benchmark(benchmark, tmpdir_factory, data_fn):

tmp_path = Path(tmpdir_factory.mktemp("dataset_ops"))

schema, data = data_fn(num_bytes=1 * GiB)

benchmark(do_write_dataset, data, tmp_path, schema)

@pytest.mark.benchmark(group="create_dataset")

def test_captioned_image(tmpdir_factory, benchmark):

write_dataset_benchmark(benchmark, tmpdir_factory, create_captioned_image_data)

@pytest.mark.benchmark(group="create_dataset")

def test_scalar(tmpdir_factory, benchmark):

write_dataset_benchmark(benchmark, tmpdir_factory, create_scalar_data)