# SPDX-License-Identifier: Apache-2.0

# SPDX-FileCopyrightText: Copyright The Lance Authors

import lance

import numpy as np

import pyarrow as pa

import pytest

pytest.importorskip("torch")

from lance.torch.data import SafeLanceDataset, get_safe_loader # noqa: E402

@pytest.fixture(scope="module")

def temp_lance_dataset(tmp_path_factory):

"""Create temporary Lance dataset for testing"""

test_dir = tmp_path_factory.mktemp("lance_data")

dataset_path = test_dir / "test_dataset.lance"

# Generate test data with batch_size aligned sample count

num_samples = 96 # 16 samples/batch * 6 batches

data = pa.table(

{

"id": range(num_samples),

"embedding": [

np.random.rand(128).astype(np.float32).tobytes()

for _ in range(num_samples)

],

}

)

lance.write_dataset(data, dataset_path)

yield str(dataset_path)

def test_dataset_initialization(temp_lance_dataset):

"""Verify dataset basic functionality"""

ds = SafeLanceDataset(temp_lance_dataset)

# Validate metadata

assert len(ds) == 96, "Sample count should match configured size"

# Validate single sample format

sample = ds[0]

assert isinstance(sample, dict), "Sample should be dictionary type"

assert {"id", "embedding"}.issubset(sample.keys()), "Missing required fields"

def test_multiprocess_loading(temp_lance_dataset, capsys):

"""Verify multi-worker data loading"""

dataset = SafeLanceDataset(temp_lance_dataset)

loader = get_safe_loader(

dataset,

num_workers=2,

batch_size=16,

drop_last=False, # Ensure full batches

)

total_samples = 0

for batch in loader:

assert batch["id"].shape == (16,), "Batch dimension mismatch"

total_samples += batch["id"].shape[0]

# Validate complete dataset loading

assert total_samples == 96, "Should load all samples"