# SPDX-License-Identifier: Apache-2.0

# SPDX-FileCopyrightText: Copyright The Lance Authors

from pathlib import Path

import lance

import numpy as np

import pyarrow as pa

import pytest

datasets = pytest.importorskip("datasets")

pil = pytest.importorskip("PIL")

def test_write_hf_dataset(tmp_path: Path):

hf_ds = datasets.load_dataset(

"cornell-movie-review-data/rotten_tomatoes",

split="train[:50]",

)

ds = lance.write_dataset(hf_ds, tmp_path)

assert ds.count_rows() == 50

assert ds.schema == hf_ds.features.arrow_schema

@pytest.mark.cuda

def test_image_hf_dataset(tmp_path: Path):

import lance.torch.data

ds = datasets.Dataset.from_dict(

{"i": [np.zeros(shape=(16, 16, 3), dtype=np.uint8)]},

features=datasets.Features({"i": datasets.Image()}),

)

ds = lance.write_dataset(ds, tmp_path)

dataset = lance.torch.data.LanceDataset(

ds,

columns=["i"],

batch_size=8,

)

batch = next(iter(dataset))

assert len(batch) == 1

assert all(

(isinstance(img, pil.Image.Image) and np.all(np.array(img) == 0))

for img in batch

)

def test_iterable_dataset(tmp_path: Path):

# IterableDataset yields dict of arrays

def gen():

yield {"text": "Good", "label": 0}

yield {"text": "Bad", "label": 1}

arrow_schema = pa.schema([("text", pa.string()), ("label", pa.int64())])

features = datasets.Features.from_arrow_schema(arrow_schema)

iter_ds = datasets.IterableDataset.from_generator(gen, features=features)

# streaming batch size is controlled by max_rows_per_group

ds1 = lance.write_dataset(iter_ds, tmp_path / "ds1.lance")

assert ds1.count_rows() == 2

assert ds1.schema == iter_ds.features.arrow_schema

# to manually control streaming batch size

ds2 = lance.write_dataset(

pa.Table.from_arrays([[], []], schema=arrow_schema), tmp_path / "ds2.lance"

)

for batch in iter_ds.iter(batch_size=1):

# shouldn't fail

ds2 = lance.write_dataset(batch, tmp_path / "ds2.lance", mode="append")

assert len(ds1) == len(ds2)

assert ds1.schema == ds2.schema