# 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.file import stable_version

def prep_dataset(tmp_path: Path, version, name: str):

dataset_dir = (

Path(__file__).parent.parent.parent.parent / "test_data" / version / name

)

shutil.copytree(dataset_dir, tmp_path / version / name)

return lance.dataset(tmp_path / version / name)

def test_add_data_storage_version(tmp_path: Path):

"""

In version 0.15 and below we did not have a data storage version. We had

writer flags that were used to determine if we should use the old or new

storage format. In version 0.16 we added a data storage version to the

manifest that should be correctly populated from existing files and/or the

writer flags

"""

tab = pa.table({"x": range(1024)})

def check_dataset(dataset_name: str, expected_version: str):

ds = prep_dataset(tmp_path, "v0.15.0", dataset_name)

assert ds.data_storage_version == expected_version

lance.write_dataset(tab, ds.uri, mode="append")

assert ds.data_storage_version == expected_version

check_dataset("v1_no_files.lance", "0.1")

check_dataset("v2_no_files.lance", stable_version())

check_dataset("v1_with_files.lance", "0.1")

check_dataset("v2_with_files.lance", "2.0")

def test_fix_data_storage_version(tmp_path: Path):

"""

In versions above 0.15 and below 0.17 we wrote the data storage version but

we might have written it incorrectly. In version 0.17 we fixed the inference

rules. However, this means a dataset could exist that has v2 files and a

data storage version of 0.1. Or, worse, it is even possible to have created

a dataset with a mix of 0.1 and 2.0 files. The former, we can fix automatically

and the latter we can at least detect and advise the user rollback their dataset.

"""

ds = prep_dataset(tmp_path, "v0.16.0", "wrong_data_version_can_fix.lance")

assert ds.data_storage_version == "0.1"

ds.delete("false")

assert ds.data_storage_version == "2.0"

ds = prep_dataset(tmp_path, "v0.16.0", "wrong_data_version_no_fix.lance")

assert ds.data_storage_version == "0.1"

with pytest.raises(

OSError, match="The dataset contains a mixture of file versions"

):

ds.delete("false")

def test_old_btree_bitmap_indices(tmp_path: Path):

"""

In versions below 0.21.0 we used the legacy file format for btree and bitmap

indices. In version 0.21.0 we switched to the new format. This test ensures

that we can still read the old indices.

"""

ds = prep_dataset(tmp_path, "v0.20.0", "old_btree_bitmap_indices.lance")

def query(filt: str):

table = ds.to_table(filter=filt)

assert table == pa.table({"bitmap": [3, 4], "btree": [3, 4]})

explain = ds.scanner(filter=filt).explain_plan()

assert "ScalarIndexQuery" in explain or "MaterializeIndex" in explain

query("bitmap > 2")

query("btree > 2")

def test_index_no_details(tmp_path: Path):

"""

In versions below 0.19.3 we did not write index details to the index metadata.

This test ensures that we can still read the old indices.

"""

ds = prep_dataset(tmp_path, "v0.18.2", "index_no_details.lance")

assert ds.to_table(filter="id > 2").num_rows == 97