# SPDX-License-Identifier: Apache-2.0

# SPDX-FileCopyrightText: Copyright The Lance Authors

import base64

import contextlib

import os

import pickle

import platform

import random

import re

import time

import uuid

from datetime import date, datetime, timedelta

from pathlib import Path

from typing import List

from unittest import mock

import lance

import lance.fragment

import numpy as np

import pandas as pd

import pandas.testing as tm

import polars as pl

import pyarrow as pa

import pyarrow.compute as pc

import pyarrow.dataset as pa_ds

import pyarrow.parquet as pq

import pytest

from helper import ProgressForTest

from lance._dataset.sharded_batch_iterator import ShardedBatchIterator

from lance.commit import CommitConflictError

from lance.dataset import LANCE_COMMIT_MESSAGE_KEY, AutoCleanupConfig

from lance.debug import format_fragment

from lance.file import stable_version

from lance.schema import LanceSchema

from lance.util import validate_vector_index

# Various valid inputs for write_dataset

input_schema = pa.schema([pa.field("a", pa.float64()), pa.field("b", pa.int64())])

input_data = [

# (schema, data)

(None, pa.table({"a": [1.0, 2.0], "b": [20, 30]})),

(None, pa.record_batch([[1.0, 2.0], [20, 30]], names=["a", "b"])),

(None, pd.DataFrame({"a": [1.0, 2.0], "b": [20, 30]})),

(None, pl.DataFrame({"a": [1.0, 2.0], "b": [20, 30]})),

(

input_schema,

[pa.record_batch([pa.array([1.0, 2.0]), pa.array([20, 30])], names=["a", "b"])],

),

# Can provide an iterator with a schema that is different but cast-able

(

input_schema,

iter(

pa.table(

{

"a": [1.0, 2.0],

"b": pa.array([20, 30], pa.int32()),

}

).to_batches()

),

),

]

@pytest.mark.parametrize("schema,data", input_data, ids=type)

def test_input_data(tmp_path: Path, schema, data):

base_dir = tmp_path / "test"

dataset = lance.write_dataset(data, base_dir, schema=schema)

assert dataset.to_table() == input_data[0][1]

def test_roundtrip_types(tmp_path: Path):

table = pa.table(

{

"dict": pa.array(["a", "b", "a"], pa.dictionary(pa.int8(), pa.string())),

# PyArrow doesn't support creating large_string dictionaries easily.

"large_dict": pa.DictionaryArray.from_arrays(

pa.array([0, 1, 1], pa.int8()),

pa.array(["foo", "bar"], pa.large_string()),

),

"list": pa.array(

[["a", "b"], ["c", "d"], ["e", "f"]], pa.list_(pa.string())

),

"large_list": pa.array(

[["a", "b"], ["c", "d"], ["e", "f"]], pa.large_list(pa.string())

),

}

)

# TODO: V2 does not currently handle large_dict

dataset = lance.write_dataset(table, tmp_path, data_storage_version="legacy")

assert dataset.schema == table.schema

assert dataset.to_table() == table

def test_dataset_overwrite(tmp_path: Path):

table1 = pa.Table.from_pylist([{"a": 1, "b": 2}, {"a": 10, "b": 20}])

base_dir = tmp_path / "test"

lance.write_dataset(table1, base_dir)

table2 = pa.Table.from_pylist([{"s": "one"}, {"s": "two"}])

lance.write_dataset(table2, base_dir, mode="overwrite")

dataset = lance.dataset(base_dir)

assert dataset.to_table() == table2

ds_v1 = lance.dataset(base_dir, version=1)

assert ds_v1.to_table() == table1

def test_truncate_table(tmp_path: Path):

base_dir = tmp_path / "truncate"

table = pa.table(

{

"i": pa.array([1, 2, 3], pa.int32()),

"dict": pa.DictionaryArray.from_arrays(

pa.array([0, 1, 2], pa.uint16()), pa.array(["a", "b", "c"])

),

}

)

ds = lance.write_dataset(table, base_dir, data_storage_version="stable")

assert ds.count_rows() == 3

ds.truncate_table()

assert ds.count_rows() == 0

assert ds.schema == table.schema

def test_dataset_append(tmp_path: Path):

table = pa.Table.from_pydict({"colA": [1, 2, 3], "colB": [4, 5, 6]})

base_dir = tmp_path / "test"

# verify append works even if no dataset existed at the uri

lance.write_dataset(table, base_dir, mode="append")

dataset = lance.dataset(base_dir)

assert dataset.to_table() == table

# verify appending batches with a different schema doesn't work

table2 = pa.Table.from_pydict({"COLUMN-C": [1, 2, 3], "colB": [4, 5, 6]})

with pytest.raises(OSError):

lance.write_dataset(table2, dataset, mode="append")

# But we can append subschemas

table3 = pa.Table.from_pydict({"colA": [4, 5, 6]})

dataset.insert(table3) # Append is default

assert dataset.to_table() == pa.table(

{"colA": [1, 2, 3, 4, 5, 6], "colB": [4, 5, 6, None, None, None]}

)

def test_dataset_from_record_batch_iterable(tmp_path: Path):

base_dir = tmp_path / "test"

test_pylist = [{"colA": "Alice", "colB": 20}, {"colA": "Blob", "colB": 30}]

# split into two batches

batches = [

pa.RecordBatch.from_pylist([test_pylist[0]]),

pa.RecordBatch.from_pylist([test_pylist[1]]),

]

# define schema

schema = pa.schema(

[

pa.field("colA", pa.string()),

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

]

)

# write dataset with iterator

lance.write_dataset(iter(batches), base_dir, schema)

dataset = lance.dataset(base_dir)

# After combined into one batch, make sure it is the same as original pylist

assert list(dataset.to_batches())[0].to_pylist() == test_pylist

# write dataset with list

lance.write_dataset(batches, base_dir, schema, mode="overwrite")

# After combined into one batch, make sure it is the same as original pylist

assert list(dataset.to_batches())[0].to_pylist() == test_pylist

def test_to_batches_with_partial_last_batch(tmp_path: Path):

base_dir = tmp_path / "test_batches"

row_count_per_file = 32

batch_size = 5

# Generate 3 batches of 32 rows each (96 total)

pylist = [{"colA": f"Row{i}", "colB": i} for i in range(row_count_per_file * 3)]

batches = [

pa.RecordBatch.from_pylist(

pylist[i * row_count_per_file : (i + 1) * row_count_per_file]

)

for i in range(3)

]

# Write dataset

schema = pa.schema([pa.field("colA", pa.string()), pa.field("colB", pa.int64())])

lance.write_dataset(batches, base_dir, schema, max_rows_per_file=row_count_per_file)

dataset = lance.dataset(base_dir)

# Check batch sizes

# strict_batch_size = True, batch_size = 5(batch_size < row_count_per_file),

all_batches = list(

dataset.to_batches(batch_size=batch_size, strict_batch_size=True)

)

print(all_batches)

assert sum(b.num_rows for b in all_batches) == row_count_per_file * 3 # Total rows

assert all(b.num_rows == batch_size for b in all_batches[:-1]) # Full batches

assert all_batches[-1].num_rows == 1 # Final partial batch

# Verify data integrity

combined = [row for batch in all_batches for row in batch.to_pylist()]

assert combined == pylist

# strict_batch_size = True, batch_size = 5*10 (batch_size > row_count_per_file),

large_batch_size = 10 * batch_size

all_batches = list(

dataset.to_batches(batch_size=large_batch_size, strict_batch_size=True)

)

assert sum(b.num_rows for b in all_batches) == row_count_per_file * 3 # Total rows

assert all(b.num_rows == large_batch_size for b in all_batches[:-1]) # Full batches

assert all_batches[-1].num_rows == 46 # Final partial batch

# strict_batch_size = False

# fragment 32 rows --> [5,5,5,5,5,5,2]

all_batches = list(

dataset.to_batches(batch_size=batch_size, strict_batch_size=False)

)

assert sum(b.num_rows for b in all_batches) == row_count_per_file * 3 # Total rows

partial_batches = [b for b in all_batches if b.num_rows < batch_size]

partial_batch_count = 3

assert len(partial_batches) == partial_batch_count

assert all(b.num_rows == 2 for b in partial_batches[:-1])

# strict_batch_size = False, batch_size = 5*10 (batch_size > row_count_per_file),

all_batches = list(

dataset.to_batches(batch_size=large_batch_size, strict_batch_size=False)

)

assert sum(b.num_rows for b in all_batches) == row_count_per_file * 3 # Total rows

assert all(b.num_rows == 32 for b in all_batches) # Full batches

# 32 rows, 1 row per file

ds = lance.write_dataset(

pa.table({"a": range(32)}), base_dir, max_rows_per_file=1, mode="overwrite"

)

# 1 row per batch if strict_batch_size is False (regardless of batch_size)

for batch in ds.to_batches():

assert batch.num_rows == 1

for batch in ds.to_batches(batch_size=8):

assert batch.num_rows == 1

# We should get 8 rows per batch if strict_batch_size is True

for batch in ds.to_batches(batch_size=8, strict_batch_size=True):

assert batch.num_rows == 8

def test_schema_metadata(tmp_path: Path):

schema = pa.schema(

[

pa.field("a", pa.int64(), metadata={b"thisis": "a"}),

pa.field("b", pa.int64(), metadata={b"thisis": "b"}),

],

metadata={b"foo": b"bar", b"baz": b"qux"},

)

table = pa.Table.from_pydict({"a": range(100), "b": range(100)}, schema=schema)

ds = lance.write_dataset(table, tmp_path)

# Original schema

assert ds.schema.metadata == {b"foo": b"bar", b"baz": b"qux"}

assert ds.schema.field("a").metadata == {b"thisis": b"a"}

assert ds.schema.field("b").metadata == {b"thisis": b"b"}

# Replace schema metadata (using new unified API)

ds.update_schema_metadata({"foo": "baz"}, replace=True)

assert ds.schema.metadata == {b"foo": b"baz"}

assert ds.schema.field("a").metadata == {b"thisis": b"a"}

assert ds.schema.field("b").metadata == {b"thisis": b"b"}

# Replace field metadata (using new unified API)

# Use field path instead of field ID

ds.update_field_metadata({"a": {"thisis": "c"}}, replace=True)

assert ds.schema.field("a").metadata == {b"thisis": b"c"}

assert ds.schema.field("b").metadata == {b"thisis": b"b"}

# Overwrite overwrites metadata

ds = lance.write_dataset(table, tmp_path, mode="overwrite")

assert ds.schema.metadata == {b"foo": b"bar", b"baz": b"qux"}

assert ds.schema.field("a").metadata == {b"thisis": b"a"}

assert ds.schema.field("b").metadata == {b"thisis": b"b"}

def test_versions(tmp_path: Path):

table1 = pa.Table.from_pylist([{"a": 1, "b": 2}, {"a": 10, "b": 20}])

base_dir = tmp_path / "test"

lance.write_dataset(table1, base_dir)

assert len(lance.dataset(base_dir).versions()) == 1

table2 = pa.Table.from_pylist([{"s": "one"}, {"s": "two"}])

time.sleep(1)

lance.write_dataset(table2, base_dir, mode="overwrite")

assert len(lance.dataset(base_dir).versions()) == 2

v1, v2 = lance.dataset(base_dir).versions()

assert v1["version"] == 1

assert v2["version"] == 2

assert isinstance(v1["timestamp"], datetime)

assert isinstance(v2["timestamp"], datetime)

assert v1["timestamp"] < v2["timestamp"]

assert isinstance(v1["metadata"], dict)

assert isinstance(v2["metadata"], dict)

def test_version_id(tmp_path: Path):

table1 = pa.Table.from_pylist([{"a": 1, "b": 2}, {"a": 10, "b": 20}])

base_dir = tmp_path / "test"

original_ds = lance.write_dataset(table1, base_dir)

assert original_ds.version == 1

assert original_ds.latest_version == 1

table2 = pa.Table.from_pylist([{"s": "one"}, {"s": "two"}])

time.sleep(1)

updated_ds = lance.write_dataset(table2, base_dir, mode="overwrite")

assert original_ds.version == 1

assert original_ds.latest_version == 2

assert updated_ds.version == 2

assert updated_ds.latest_version == 2

def test_checkout(tmp_path: Path):

tab = pa.table({"a": range(3)})

ds1 = lance.write_dataset(tab, tmp_path)

ds1.delete("a = 1")

ds2 = ds1.checkout_version(1)

assert ds2.version == 1

assert ds2.to_table() == tab

assert ds1.version == 2

assert ds1.to_table() == pa.table({"a": [0, 2]})

ds1.delete("a = 2")

assert ds1.count_rows() == 1

assert ds2.checkout_version(ds2.latest_version).version == ds1.version

def test_asof_checkout(tmp_path: Path):

table = pa.Table.from_pydict({"colA": [1, 2, 3], "colB": [4, 5, 6]})

base_dir = tmp_path / "test"

lance.write_dataset(table, base_dir)

assert len(lance.dataset(base_dir).versions()) == 1

time.sleep(0.1)

ts_1 = datetime.now()

time.sleep(0.1)

lance.write_dataset(table, base_dir, mode="append")

assert len(lance.dataset(base_dir).versions()) == 2

time.sleep(0.1)

ts_2 = datetime.now()

time.sleep(0.1)

lance.write_dataset(table, base_dir, mode="append")

assert len(lance.dataset(base_dir).versions()) == 3

time.sleep(0.1)

ts_3 = datetime.now()

# check that only the first batch is present

ds = lance.dataset(base_dir, asof=ts_1)

assert ds.version == 1

assert len(ds.to_table()) == 3

# check that the first and second batch are present

ds = lance.dataset(base_dir, asof=ts_2)

assert ds.version == 2

assert len(ds.to_table()) == 6

# check that all batches are present

ds = lance.dataset(base_dir, asof=ts_3)

assert ds.version == 3

assert len(ds.to_table()) == 9

def test_enable_stable_row_ids(tmp_path: Path):

table = pa.Table.from_pylist(

[{"name": "Alice", "age": 20}, {"name": "Bob", "age": 30}]

)

lance.write_dataset(table, tmp_path, enable_stable_row_ids=True)

ds = lance.write_dataset(table, tmp_path, enable_stable_row_ids=True, mode="append")

table_before = ds.scanner(with_row_id=True, with_row_address=True).to_table()

assert len(table_before) == 4

assert table_before["_rowid"][0].as_py() == 0

assert table_before["_rowid"][1].as_py() == 1

assert table_before["_rowid"][2].as_py() == 2

assert table_before["_rowid"][3].as_py() == 3

assert table_before["_rowaddr"][0].as_py() == 0

assert table_before["_rowaddr"][1].as_py() == 1

assert table_before["_rowaddr"][2].as_py() == (1 << 32) + 0

assert table_before["_rowaddr"][3].as_py() == (1 << 32) + 1

ds.optimize.compact_files()

table_after = ds.scanner(with_row_id=True, with_row_address=True).to_table()

assert len(table_after) == 4

assert table_after["_rowid"][0].as_py() == 0

assert table_after["_rowid"][1].as_py() == 1

assert table_after["_rowid"][2].as_py() == 2

assert table_after["_rowid"][3].as_py() == 3

assert table_after["_rowaddr"][0].as_py() == (2 << 32) + 0

assert table_after["_rowaddr"][1].as_py() == (2 << 32) + 1

assert table_after["_rowaddr"][2].as_py() == (2 << 32) + 2

assert table_after["_rowaddr"][3].as_py() == (2 << 32) + 3

def test_v2_manifest_paths(tmp_path: Path):

lance.write_dataset(

pa.table({"a": range(100)}), tmp_path, enable_v2_manifest_paths=True

)

manifest_path = os.listdir(tmp_path / "_versions")

assert len(manifest_path) == 1

assert re.match(r"\d{20}\.manifest", manifest_path[0])

def test_default_v2_manifest_paths(tmp_path: Path):

lance.write_dataset(pa.table({"a": range(100)}), tmp_path)

manifest_path = os.listdir(tmp_path / "_versions")

assert len(manifest_path) == 1

assert re.match(r"\d{20}\.manifest", manifest_path[0])

def test_v2_manifest_paths_migration(tmp_path: Path):

# Create a dataset with v1 manifest paths

lance.write_dataset(

pa.table({"a": range(100)}), tmp_path, enable_v2_manifest_paths=False

)

manifest_path = os.listdir(tmp_path / "_versions")

assert manifest_path == ["1.manifest"]

# Migrate to v2 manifest paths

lance.dataset(tmp_path).migrate_manifest_paths_v2()

manifest_path = os.listdir(tmp_path / "_versions")

assert len(manifest_path) == 1

assert re.match(r"\d{20}\.manifest", manifest_path[0])

def test_tag(tmp_path: Path):

table = pa.Table.from_pydict({"colA": [1, 2, 3], "colB": [4, 5, 6]})

base_dir = tmp_path / "test"

lance.write_dataset(table, base_dir)

ds = lance.write_dataset(table, base_dir, mode="append")

assert len(ds.tags.list()) == 0

with pytest.raises(ValueError):

ds.tags.create("tag1", 3)

with pytest.raises(ValueError):

ds.tags.delete("tag1")

ds.tags.create("tag1", 1)

assert len(ds.tags.list()) == 1

with pytest.raises(ValueError):

ds.tags.create("tag1", 1)

ds.tags.delete("tag1")

ds.tags.create("tag1", 1)

ds.tags.create("tag2", 1)

assert len(ds.tags.list()) == 2

with pytest.raises(OSError):

ds.checkout_version("tag3")

assert ds.checkout_version("tag1").version == 1

ds = lance.dataset(base_dir, "tag1")

assert ds.version == 1

with pytest.raises(ValueError):

lance.dataset(base_dir, "missing-tag")

# test tag update

with pytest.raises(

ValueError, match="Version not found error: version main:3 does not exist"

):

ds.tags.update("tag1", 3)

with pytest.raises(

ValueError, match="Ref not found error: tag tag3 does not exist"

):

ds.tags.update("tag3", 1)

ds.tags.update("tag1", 2)

ds = lance.dataset(base_dir, "tag1")

assert ds.version == 2

ds.tags.update("tag1", 1)

ds = lance.dataset(base_dir, "tag1")

assert ds.version == 1

version = ds.tags.get_version("tag1")

assert version == 1

ds.create_branch("branch", "tag1")

ds.tags.create("tag3", ("branch", None))

target_tag = ds.tags.list().get("tag3")

assert ds.tags.get_version("tag3") == 1

assert len(ds.tags.list()) == 3

assert target_tag is not None

assert target_tag["version"] == 1

assert target_tag["branch"] == "branch"

ds.tags.update("tag3", (None, 2))

target_tag = ds.tags.list()["tag3"]

assert ds.tags.get_version("tag3") == 2

assert target_tag is not None

assert target_tag["version"] == 2

assert target_tag["branch"] is None

ds.create_branch("branch2", 2)

ds.tags.update("tag3", ("branch2", 2))

target_tag = ds.tags.list()["tag3"]

assert ds.tags.get_version("tag3") == 2

assert target_tag is not None

assert target_tag["version"] == 2

assert target_tag["branch"] == "branch2"

ds.tags.delete("tag3")

assert len(ds.tags.list()) == 2

def test_tag_order(tmp_path: Path):

table = pa.Table.from_pydict({"colA": [1, 2, 3], "colB": [4, 5, 6]})

base_dir = tmp_path / "test"

for i in range(3):

mode = "append" if i > 0 else "create"

ds = lance.write_dataset(table, base_dir, mode=mode)

expected_tags = {"tag3": 3, "tag2": 2, "tag1": 1}

for name, version in expected_tags.items():

ds.tags.create(name, version)

tags_asc = ds.tags.list_ordered(order="asc")

assert len(tags_asc) == 3

tag_names_asc = [t[0] for t in tags_asc]

assert tag_names_asc == sorted(expected_tags.keys()), (

f"Unexpected ascending order: {tag_names_asc}"

)

# Test descending order (default)

tags_desc = ds.tags.list_ordered(order="desc")

assert len(tags_desc) == 3

tag_names_desc = [t[0] for t in tags_desc]

assert tag_names_desc == list(expected_tags.keys()), (

f"Unexpected descending order: {tag_names_desc}"

)

# Test without parameter (should default to descending)

tags_default = ds.tags.list_ordered()

assert tags_default == tags_desc, "Default order should match descending order"

def test_sample(tmp_path: Path):

table1 = pa.Table.from_pydict({"x": [0, 10, 20, 30, 40, 50], "y": range(6)})

base_dir = tmp_path / "test"

lance.write_dataset(table1, base_dir)

dataset = lance.dataset(base_dir)

sampled = dataset.sample(3)

assert sampled.num_rows == 3

assert sampled.num_columns == 2

sampled = dataset.sample(4, columns=["x"])

assert sampled.num_rows == 4

assert sampled.num_columns == 1

for row in sampled.column(0).chunk(0):

assert row.as_py() % 10 == 0

def test_take(tmp_path: Path):

table1 = pa.Table.from_pylist([{"a": 1, "b": 2}, {"a": 10, "b": 20}])

base_dir = tmp_path / "test"

lance.write_dataset(table1, base_dir)

dataset = lance.dataset(base_dir)

table2 = dataset.take([0, 1])

assert isinstance(table2, pa.Table)

assert table2 == table1

def test_take_rowid_rowaddr(tmp_path: Path):

sample_size = 10

table1 = pa.table({"a": range(1000), "b": range(1000)})

base_dir = tmp_path / "test_take_rowid_rowaddr"

lance.write_dataset(

table1, base_dir, enable_stable_row_ids=False, max_rows_per_file=50

)

dataset = lance.dataset(base_dir)

total_rows = len(dataset)

sampled_indices = random.sample(range(total_rows), min(sample_size, total_rows))

sample_dataset = dataset.take(sampled_indices, columns=["_rowid"])

assert sample_dataset.num_rows == 10

assert sample_dataset.num_columns == 1

sample_dataset = dataset.take(sampled_indices, columns=["_rowid", "_rowid"])

assert sample_dataset.num_rows == 10

assert sample_dataset.num_columns == 2

sample_dataset = dataset.take([1, 2, 3, 4, 5, 6, 7, 8, 9, 100], columns=["_rowid"])

assert sample_dataset.num_rows == 10

assert sample_dataset.num_columns == 1

sample_dataset = dataset.take(sampled_indices, columns=["_rowaddr"])

assert sample_dataset.num_rows == 10

assert sample_dataset.num_columns == 1

sample_dataset = dataset.take(

[1, 2, 3, 4, 5, 6, 7, 8, 9, 100], columns=["_rowaddr"]

)

assert sample_dataset.num_rows == 10

assert sample_dataset.num_columns == 1

sample_dataset = dataset.take(sampled_indices, columns=["_rowaddr", "_rowid"])

assert sample_dataset.num_rows == 10

assert sample_dataset.num_columns == 2

sample_dataset = dataset.take(

[1, 2, 3, 4, 5, 6, 7, 8, 9, 100], columns=["_rowaddr", "_rowid"]

)

assert sample_dataset.num_rows == 10

assert sample_dataset.num_columns == 2

sample_dataset = dataset.take(sampled_indices, columns=["_rowid", "_rowaddr"])

assert sample_dataset.num_rows == 10

assert sample_dataset.num_columns == 2

sample_dataset = dataset.take(

[1, 2, 3, 4, 5, 6, 7, 8, 9, 100], columns=["_rowid", "_rowaddr"]

)

assert sample_dataset.num_rows == 10

assert sample_dataset.num_columns == 2

sample_dataset = dataset.take(sampled_indices, columns=["a", "_rowid", "_rowaddr"])

assert sample_dataset.num_rows == 10

assert sample_dataset.num_columns == 3

sample_dataset = dataset.take(

[1, 2, 3, 4, 5, 6, 7, 8, 9, 100], columns=["a", "_rowid", "_rowaddr"]

)

assert sample_dataset.num_rows == 10

assert sample_dataset.num_columns == 3

sample_dataset = dataset.take(sampled_indices, columns=["_rowid", "_rowaddr", "b"])

assert sample_dataset.num_rows == 10

assert sample_dataset.num_columns == 3

sample_dataset = dataset.take(

[1, 2, 3, 4, 5, 6, 7, 8, 9, 100], columns=["_rowid", "_rowaddr", "b"]

)

assert sample_dataset.num_rows == 10

assert sample_dataset.num_columns == 3

sample_dataset = dataset.take(sampled_indices, columns=["a", "b"])

assert sample_dataset.num_rows == 10

assert sample_dataset.num_columns == 2

sample_dataset = dataset.take([1, 2, 3, 4, 5, 6, 7, 8, 9, 100], columns=["a", "b"])

assert sample_dataset.num_rows == 10

assert sample_dataset.num_columns == 2

@pytest.mark.parametrize(

"column_name",

[

"_rowid",

"_rowaddr",

"_rowoffset",

"_row_created_at_version",

"_row_last_updated_at_version",

],

)

def test_take_system_columns_values(tmp_path: Path, column_name: str):

"""Test that system columns return correct values in take."""

table = pa.table({"a": range(100), "b": range(100, 200)})

base_dir = tmp_path / "test_take_system_columns_values"

# Use max_rows_per_file to create multiple fragments

lance.write_dataset(table, base_dir, max_rows_per_file=25)

dataset = lance.dataset(base_dir)

indices = [0, 5, 10, 50, 99]

result = dataset.take(indices, columns=[column_name, "a"])

assert result.num_rows == len(indices)

assert result.schema.names == [column_name, "a"]

col_values = result.column(column_name).to_pylist()

a_values = result.column("a").to_pylist()

# Verify column type is UInt64

assert result.column(column_name).type == pa.uint64()

# Verify data column values

assert a_values == indices

# Verify system column values based on column type

if column_name == "_rowid":

# Without stable row IDs, _rowid equals _rowaddr (not the index).

# Row address = (fragment_id << 32) | row_offset_within_fragment

# With max_rows_per_file=25: frag0=0-24, frag1=25-49, frag2=50-74, frag3=75-99

expected_rowids = [

(0 << 32) | 0, # index 0: fragment 0, offset 0

(0 << 32) | 5, # index 5: fragment 0, offset 5

(0 << 32) | 10, # index 10: fragment 0, offset 10

(2 << 32) | 0, # index 50: fragment 2, offset 0

(3 << 32) | 24, # index 99: fragment 3, offset 24

]

assert col_values == expected_rowids

elif column_name in ("_row_created_at_version", "_row_last_updated_at_version"):

# All rows created/updated at version 1

assert col_values == [1] * len(indices)

# _rowaddr and _rowoffset values depend on fragment layout

def test_take_system_columns_column_ordering(tmp_path: Path):

"""Test that column ordering is preserved when using system columns."""

table = pa.table({"a": range(50), "b": range(50, 100)})

base_dir = tmp_path / "test_take_column_ordering"

lance.write_dataset(table, base_dir)

dataset = lance.dataset(base_dir)

indices = [0, 1, 2]

# Test different orderings with all system columns

result = dataset.take(indices, columns=["_rowid", "a", "_rowaddr"])

assert result.schema.names == ["_rowid", "a", "_rowaddr"]

result = dataset.take(indices, columns=["a", "_rowaddr", "_rowid"])

assert result.schema.names == ["a", "_rowaddr", "_rowid"]

result = dataset.take(indices, columns=["_rowaddr", "_rowid", "b", "a"])

assert result.schema.names == ["_rowaddr", "_rowid", "b", "a"]

# Test with version columns

result = dataset.take(

indices,

columns=[

"_row_created_at_version",

"a",

"_row_last_updated_at_version",

"_rowid",

],

)

assert result.schema.names == [

"_row_created_at_version",

"a",

"_row_last_updated_at_version",

"_rowid",

]

# Test with all system columns in mixed order

result = dataset.take(

indices,

columns=[

"_rowoffset",

"_row_last_updated_at_version",

"b",

"_rowaddr",

"_row_created_at_version",

"a",

"_rowid",

],

)

assert result.schema.names == [

"_rowoffset",

"_row_last_updated_at_version",

"b",

"_rowaddr",

"_row_created_at_version",

"a",

"_rowid",

]

def test_take_version_system_columns(tmp_path: Path):

"""Test _row_created_at_version and _row_last_updated_at_version columns."""

table = pa.table({"a": range(50)})

base_dir = tmp_path / "test_take_version_columns"

lance.write_dataset(table, base_dir, enable_stable_row_ids=True)

dataset = lance.dataset(base_dir)

# Initial version is 1

initial_version = dataset.version

indices = [0, 10, 25]

result = dataset.take(

indices,

columns=["a", "_row_created_at_version", "_row_last_updated_at_version"],

)

assert result.num_rows == 3

created_at = result.column("_row_created_at_version").to_pylist()

updated_at = result.column("_row_last_updated_at_version").to_pylist()

# All rows were created and last updated at the initial version

assert created_at == [initial_version] * 3

assert updated_at == [initial_version] * 3

# Now update some rows by overwriting

table2 = pa.table({"a": range(50, 100)})

lance.write_dataset(table2, base_dir, mode="append")

dataset = lance.dataset(base_dir)

# New rows should have version 2

result = dataset.take([50, 60], columns=["_row_created_at_version"])

created_at = result.column("_row_created_at_version").to_pylist()

assert created_at == [dataset.version] * 2

@pytest.mark.parametrize("indices", [[], [1, 1], [1, 1, 20, 20, 21], [21, 0, 21, 1, 0]])

def test_take_duplicate_index(tmp_path: Path, indices: List[int]):

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

dataset = lance.write_dataset(

table, tmp_path, max_rows_per_group=3, max_rows_per_file=9

)

expected = table.take(pa.array(indices, pa.int64()))

assert dataset.take(indices) == expected

def test_take_with_columns(tmp_path: Path):

table1 = pa.Table.from_pylist([{"a": 1, "b": 2}, {"a": 10, "b": 20}])

base_dir = tmp_path / "test"

lance.write_dataset(table1, base_dir)

dataset = lance.dataset(base_dir)

table2 = dataset.take([0], columns=["b"])

assert table2 == pa.Table.from_pylist([{"b": 2}])

def test_take_with_projection(tmp_path: Path):

table1 = pa.Table.from_pylist([{"a": 1, "b": "x"}, {"a": 2, "b": "y"}])

base_dir = tmp_path / "test"

lance.write_dataset(table1, base_dir)

dataset = lance.dataset(base_dir)

table2 = dataset.take([0], columns={"a2": "a*2", "bup": "UPPER(b)"})

assert table2 == pa.Table.from_pylist([{"a2": 2, "bup": "X"}])

table3 = dataset._take_rows([0], columns={"a2": "a*2", "bup": "UPPER(b)"})

assert table3 == table2

def test_filter(tmp_path: Path):

table = pa.Table.from_pydict({"a": range(100), "b": range(100)})

base_dir = tmp_path / "test"

lance.write_dataset(table, base_dir)

dataset = lance.dataset(base_dir)

actual_tab = dataset.to_table(columns=["a"], filter=(pa.compute.field("b") > 50))

assert actual_tab == pa.Table.from_pydict({"a": range(51, 100)})

def test_filter_meta_columns(tmp_path: Path):

table = pa.Table.from_pydict({"a": range(100), "b": range(100)})

base_dir = tmp_path / "test"

ds = lance.write_dataset(table, base_dir)

rowids = ds.to_table(with_row_id=True, columns=[])

some_row_id = random.sample(rowids.column(0).to_pylist(), 1)[0]

filtered = ds.to_table(filter=f"_rowid = {some_row_id}", with_row_id=True)

assert len(filtered) == 1

rowaddrs = ds.to_table(with_row_address=True, columns=[])

some_row_addr = random.sample(rowaddrs.column(0).to_pylist(), 1)[0]

filtered = ds.to_table(filter=f"_rowaddr = {some_row_addr}", with_row_address=True)

assert len(filtered) == 1

@pytest.mark.parametrize("data_storage_version", ["legacy", "stable"])

def test_limit_offset(tmp_path: Path, data_storage_version: str):

table = pa.Table.from_pydict({"a": range(100), "b": range(100)})

base_dir = tmp_path / "test"

lance.write_dataset(

table, base_dir, data_storage_version=data_storage_version, max_rows_per_file=10

)

dataset = lance.dataset(base_dir)

# test just limit

assert dataset.to_table(limit=10) == table.slice(0, 10)

assert dataset.to_table(limit=100) == table.slice(0, 100)

# test just offset

assert dataset.to_table(offset=0) == table.slice(0, 100)

assert dataset.to_table(offset=10) == table.slice(10, 90)

# test both

assert dataset.to_table(offset=10, limit=10) == table.slice(10, 10)

# Slicing in the middle of fragments

assert dataset.to_table(offset=5, limit=20) == table.slice(5, 20)

# Slicing within a single fragment

assert dataset.to_table(offset=5, limit=3) == table.slice(5, 3)

# Skipping entire fragments

assert dataset.to_table(offset=50, limit=25) == table.slice(50, 25)

# Limit past the end

assert dataset.to_table(limit=101) == table.slice(0, 100)

# Limit with offset past the end

assert dataset.to_table(offset=50, limit=51) == table.slice(50, 50)

# Offset past the end

assert dataset.to_table(offset=100) == table.slice(100, 0) # Empty table

assert dataset.to_table(offset=101) == table.slice(100, 0) # Empty table

# Offset with limit past the end

assert dataset.to_table(offset=100, limit=1) == table.slice(100, 0) # Empty table

assert dataset.to_table(offset=101, limit=1) == table.slice(100, 0) # Empty table

# Invalid limit / offset

with pytest.raises(ValueError, match="Offset must be non-negative"):

assert dataset.to_table(offset=-1, limit=10) == table.slice(50, 50)

with pytest.raises(ValueError, match="Limit must be non-negative"):

assert dataset.to_table(offset=10, limit=-1) == table.slice(50, 50)

full_ds_version = dataset.version

dataset.delete("a % 2 = 0")

filt_table = table.filter((pa.compute.bit_wise_and(pa.compute.field("a"), 1)) != 0)

assert (

dataset.to_table(offset=10).combine_chunks()

== filt_table.slice(10).combine_chunks()

)

dataset = dataset.checkout_version(full_ds_version)

dataset.restore()

dataset.delete("a > 2 AND a < 7")

filt_table = table.slice(7, 1)

assert dataset.to_table(offset=3, limit=1) == filt_table

def test_relative_paths(tmp_path: Path):

# relative paths get coerced to the full absolute path

current_dir = os.getcwd()

table = pa.Table.from_pydict({"a": range(100), "b": range(100)})

rel_uri = "test.lance"

try:

os.chdir(tmp_path)

lance.write_dataset(table, rel_uri)

# relative path works in the current dir

ds = lance.dataset(rel_uri)

assert ds.to_table() == table

finally:

os.chdir(current_dir)

# relative path doesn't work in the context of a different dir

with pytest.raises(ValueError):

ds = lance.dataset(rel_uri)

# relative path gets resolved to the right absolute path

ds = lance.dataset(tmp_path / rel_uri)

assert ds.to_table() == table

@pytest.mark.skipif(

(platform.system() == "Windows"),

reason="mocking user's home folder it not working",

)

def test_tilde_paths(tmp_path: Path):

# tilde paths get resolved to the right absolute path

tilde_uri = "~/test.lance"

table = pa.Table.from_pydict({"a": range(100), "b": range(100)})

with mock.patch.dict(

os.environ, {"HOME": str(tmp_path), "USERPROFILE": str(tmp_path)}

):

# NOTE: the resolution logic is a bit finicky

# link 1 - https://docs.rs/dirs/4.0.0/dirs/fn.home_dir.html