diff --git a/tests/test_basic_read_write.py b/tests/test_basic_read_write.py index 37436190..3fd8c880 100644 --- a/tests/test_basic_read_write.py +++ b/tests/test_basic_read_write.py @@ -3,7 +3,9 @@ import os import sys import tempfile +from collections.abc import Iterator from pathlib import Path +from typing import Any import lance import lance_ray as lr @@ -24,7 +26,7 @@ @pytest.fixture -def sample_data(): +def sample_data() -> pd.DataFrame: """Create sample data for testing.""" return pd.DataFrame( { @@ -37,14 +39,14 @@ def sample_data(): @pytest.fixture -def temp_dir(): +def temp_dir() -> Iterator[str]: """Create a temporary directory for testing.""" with tempfile.TemporaryDirectory() as temp_dir: yield temp_dir @pytest.fixture -def sample_dataset(sample_data): +def sample_dataset(sample_data: pd.DataFrame) -> Dataset: """Create a Ray Dataset from sample data.""" return ray.data.from_pandas(sample_data) @@ -52,7 +54,7 @@ def sample_dataset(sample_data): class TestWriteLance: """Test cases for write_lance function.""" - def test_write_lance_basic(self, sample_dataset, temp_dir): + def test_write_lance_basic(self, sample_dataset: Dataset, temp_dir: str) -> None: """Test basic write functionality.""" path = Path(temp_dir) / "basic_write.lance" @@ -61,7 +63,9 @@ def test_write_lance_basic(self, sample_dataset, temp_dir): assert path.exists() assert path.is_dir() - def test_write_lance_with_stable_row_ids(self, sample_dataset, temp_dir): + def test_write_lance_with_stable_row_ids( + self, sample_dataset: Dataset, temp_dir: str + ) -> None: path = Path(temp_dir) / "stable_row_ids.lance" lr.write_lance( @@ -72,29 +76,31 @@ def test_write_lance_with_stable_row_ids(self, sample_dataset, temp_dir): assert lance.dataset(str(path)).has_stable_row_ids - def test_write_lance_with_schema(self, temp_dir): + def test_write_lance_with_schema(self, temp_dir: str) -> None: """Test write with explicit schema.""" path = Path(temp_dir) / "schema_write.lance" data = pd.DataFrame({"col1": [1, 2, 3], "col2": ["a", "b", "c"]}) dataset = ray.data.from_pandas(data) - schema = pa.schema( - [pa.field("col1", pa.int64()), pa.field("col2", pa.string())] - ) + schema_fields: list[pa.Field[Any]] = [ + pa.field("col1", pa.int64()), + pa.field("col2", pa.string()), + ] + schema = pa.schema(schema_fields) lr.write_lance(dataset, str(path), schema=schema) assert path.exists() - def test_write_lance_invalid_input(self, temp_dir): + def test_write_lance_invalid_input(self, temp_dir: str) -> None: """Test error handling for invalid inputs.""" path = Path(temp_dir) / "invalid.lance" with pytest.raises((ValueError, AttributeError, TypeError)): - lr.write_lance(None, str(path)) # type: ignore + lr.write_lance(None, str(path)) # type: ignore[arg-type] - def test_write_with_pandas_map_batches(self, temp_dir): - def map_fn(row): + def test_write_with_pandas_map_batches(self, temp_dir: str) -> None: + def map_fn(row: dict[str, Any]) -> dict[str, Any]: return { "id": row["id"], "name": row["name"], @@ -103,18 +109,17 @@ def map_fn(row): "extra": None, } - def to_pd(batch: pd.DataFrame): + def to_pd(batch: pd.DataFrame) -> pd.DataFrame: return batch - schema = pa.schema( - [ - pa.field("id", pa.int64()), - pa.field("name", pa.string()), - pa.field("age", pa.int32()), - pa.field("score", pa.float64()), - pa.field("extra", pa.string()), - ] - ) + schema_fields: list[pa.Field[Any]] = [ + pa.field("id", pa.int64()), + pa.field("name", pa.string()), + pa.field("age", pa.int32()), + pa.field("score", pa.float64()), + pa.field("extra", pa.string()), + ] + schema = pa.schema(schema_fields) data = pd.DataFrame( { "id": [1, 2, 3, 4, 5], @@ -137,22 +142,21 @@ def to_pd(batch: pd.DataFrame): tbl = ds.to_table() assert set(data["name"].tolist()) == set(tbl["name"].to_pylist()) - def test_write_lance_preserves_nested_struct_fields(self, temp_dir): + def test_write_lance_preserves_nested_struct_fields(self, temp_dir: str) -> None: path = Path(temp_dir) / "nested_struct.lance" - schema = pa.schema( - [ - pa.field("id", pa.int64()), - pa.field( - "meta", - pa.struct( - [ - pa.field("userId", pa.string()), - pa.field("a.b", pa.string()), - ] - ), + schema_fields: list[pa.Field[Any]] = [ + pa.field("id", pa.int64()), + pa.field( + "meta", + pa.struct( + [ + pa.field("userId", pa.string()), + pa.field("a.b", pa.string()), + ] ), - ] - ) + ), + ] + schema = pa.schema(schema_fields) table = pa.Table.from_arrays( [ pa.array([1, 2], type=pa.int64()), @@ -190,13 +194,13 @@ class TestReadLance: """Test cases for read_lance function.""" @pytest.fixture - def lance_dataset_path(self, sample_dataset, temp_dir): + def lance_dataset_path(self, sample_dataset: Dataset, temp_dir: str) -> str: """Create a Lance dataset for reading tests.""" path = Path(temp_dir) / "test_dataset.lance" lr.write_lance(sample_dataset, str(path)) return str(path) - def test_read_lance_basic(self, lance_dataset_path): + def test_read_lance_basic(self, lance_dataset_path: str) -> None: """Test basic read functionality.""" dataset = lr.read_lance(lance_dataset_path) @@ -206,7 +210,7 @@ def test_read_lance_basic(self, lance_dataset_path): assert len(df) == 5 assert list(df.columns) == ["id", "name", "age", "score"] - def test_read_lance_with_columns(self, lance_dataset_path): + def test_read_lance_with_columns(self, lance_dataset_path: str) -> None: """Test reading specific columns.""" dataset = lr.read_lance(lance_dataset_path, columns=["id", "name"]) @@ -214,7 +218,7 @@ def test_read_lance_with_columns(self, lance_dataset_path): assert list(df.columns) == ["id", "name"] assert len(df) == 5 - def test_read_lance_with_filter(self, lance_dataset_path): + def test_read_lance_with_filter(self, lance_dataset_path: str) -> None: """Test reading with filter.""" dataset = lr.read_lance(lance_dataset_path, filter="age > 30") @@ -222,7 +226,7 @@ def test_read_lance_with_filter(self, lance_dataset_path): assert len(df) == 3 assert all(df["age"] > 30) - def test_read_lance_columns_and_filter(self, lance_dataset_path): + def test_read_lance_columns_and_filter(self, lance_dataset_path: str) -> None: """Test reading with both columns and filter.""" dataset = lr.read_lance( lance_dataset_path, columns=["name", "age"], filter="age >= 35" @@ -233,14 +237,14 @@ def test_read_lance_columns_and_filter(self, lance_dataset_path): assert len(df) == 3 assert all(df["age"] >= 35) - def test_read_lance_filter_and_count(self, lance_dataset_path): + def test_read_lance_filter_and_count(self, lance_dataset_path: str) -> None: """Test reading filter and count.""" dataset = lr.read_lance( lance_dataset_path, columns=["name", "age"], filter="age >= 35" ) assert dataset.count() == 3 - def test_read_lance_nonexistent_path(self): + def test_read_lance_nonexistent_path(self) -> None: """Test reading from non-existent path.""" with pytest.raises((FileNotFoundError, OSError, Exception)): lr.read_lance("/path/that/does/not/exist") @@ -249,7 +253,9 @@ def test_read_lance_nonexistent_path(self): class TestReadWrite: """Integration tests for read and write operations.""" - def test_write_then_read_roundtrip(self, sample_data, temp_dir): + def test_write_then_read_roundtrip( + self, sample_data: pd.DataFrame, temp_dir: str + ) -> None: """Test writing data and then reading it back.""" path = Path(temp_dir) / "roundtrip.lance" @@ -267,7 +273,7 @@ def test_write_then_read_roundtrip(self, sample_data, temp_dir): pd.testing.assert_frame_equal(original_sorted, read_sorted) - def test_append_mode(self, sample_data, temp_dir): + def test_append_mode(self, sample_data: pd.DataFrame, temp_dir: str) -> None: """Test append mode with read verification.""" path = Path(temp_dir) / "append_test.lance" @@ -293,7 +299,7 @@ def test_append_mode(self, sample_data, temp_dir): assert len(full_df) == 5 # 3 initial + 2 appended - def test_overwrite_mode(self, sample_dataset, temp_dir): + def test_overwrite_mode(self, sample_dataset: Dataset, temp_dir: str) -> None: """Test different write modes.""" path = Path(temp_dir) / "modes_test.lance" @@ -326,7 +332,9 @@ def test_overwrite_mode(self, sample_dataset, temp_dir): overwritten_df = overwritten_dataset.to_pandas() assert len(overwritten_df) == 2 # Should have 2 rows after overwrite - def test_read_lance_with_fragment_ids(self, sample_dataset, temp_dir): + def test_read_lance_with_fragment_ids( + self, sample_dataset: Dataset, temp_dir: str + ) -> None: """Test reading with fragment IDs.""" path = Path(temp_dir) / "fragment_ids_test.lance" lr.write_lance( @@ -339,7 +347,7 @@ def test_read_lance_with_fragment_ids(self, sample_dataset, temp_dir): class TestAddColumns: """Test cases for add_columns function.""" - def test_add_columns_basic(self, sample_dataset, temp_dir): + def test_add_columns_basic(self, sample_dataset: Dataset, temp_dir: str) -> None: """Test basic add columns functionality.""" path = Path(temp_dir) / "add_columns_test.lance" lr.write_lance( @@ -370,7 +378,9 @@ def double_score(x: pa.RecordBatch) -> pa.RecordBatch: class TestNamespaceReadWrite: """Test cases for read/write with DirectoryNamespace.""" - def test_write_and_read_with_directory_namespace(self, sample_data, temp_dir): + def test_write_and_read_with_directory_namespace( + self, sample_data: pd.DataFrame, temp_dir: str + ) -> None: """Test write and read using DirectoryNamespace.""" table_id = ["test_table"] @@ -398,7 +408,7 @@ def test_write_and_read_with_directory_namespace(self, sample_data, temp_dir): class TestDatasetOptions: """Test cases for dataset options in LanceDataset.""" - def test_dataset_with_version(self, sample_dataset, temp_dir): + def test_dataset_with_version(self, sample_dataset: Dataset, temp_dir: str) -> None: """Test dataset options like version and block size.""" path = Path(temp_dir) / "dataset_options_test.lance" lr.write_lance(sample_dataset, str(path)) @@ -429,12 +439,12 @@ def test_dataset_with_version(self, sample_dataset, temp_dir): from lance import DatasetBasePath, blob_array, blob_field except Exception: - class _Missing: # type: ignore[no-redef] + class _Missing: pass - DatasetBasePath = _Missing - blob_array = _Missing - blob_field = _Missing + DatasetBasePath = _Missing # type: ignore[assignment,misc] + blob_array = _Missing # type: ignore[assignment] + blob_field = _Missing # type: ignore[assignment] class TestMultiBaseLayout: @@ -457,7 +467,7 @@ class TestMultiBaseLayout: ``Duplicate base path ID 0`` error. """ - def test_multiple_initial_bases_without_explicit_id(self, temp_dir): + def test_multiple_initial_bases_without_explicit_id(self, temp_dir: str) -> None: """Multiple DatasetBasePath objects without explicit id should not collide. When the user provides two (or more) ``DatasetBasePath`` objects @@ -494,7 +504,7 @@ def test_multiple_initial_bases_without_explicit_id(self, temp_dir): ds = lance.dataset(uri) assert ds.count_rows() == 3 - base_paths = ds._ds.base_paths() + base_paths = ds._ds.base_paths() # type: ignore[attr-defined] assert len(base_paths) >= 2 base_ids = list(base_paths.keys()) assert len(set(base_ids)) == len(base_ids), ( @@ -505,7 +515,7 @@ def test_multiple_initial_bases_without_explicit_id(self, temp_dir): bool(missing_fragment_write_options("base_store_params")), reason=fragment_write_options_skip_reason("base_store_params"), ) - def test_multiple_initial_bases_with_blob_v2(self, temp_dir): + def test_multiple_initial_bases_with_blob_v2(self, temp_dir: str) -> None: """Multi-base write/read with blob v2 columns and no explicit IDs. This is the full end-to-end scenario: blob data lives across @@ -536,7 +546,7 @@ def test_multiple_initial_bases_with_blob_v2(self, temp_dir): ) ray_ds = ray.data.from_arrow(table) - base_store_params = { + base_store_params: dict[str, dict[str, Any]] = { base1_dir.as_uri(): {}, base2_dir.as_uri(): {}, } @@ -557,14 +567,14 @@ def test_multiple_initial_bases_with_blob_v2(self, temp_dir): ds = lance.dataset(uri, base_store_params=base_store_params) assert ds.count_rows() == 2 - base_paths = ds._ds.base_paths() + base_paths = ds._ds.base_paths() # type: ignore[attr-defined] assert len(base_paths) >= 2 base_ids = list(base_paths.keys()) assert len(set(base_ids)) == len(base_ids), ( f"Base path IDs must be unique, got: {base_paths}" ) - def test_explicit_ids_are_preserved(self, temp_dir): + def test_explicit_ids_are_preserved(self, temp_dir: str) -> None: """When the user provides explicit IDs, they must be preserved.""" base1_dir = Path(temp_dir) / "explicit_base1" base2_dir = Path(temp_dir) / "explicit_base2" @@ -594,11 +604,11 @@ def test_explicit_ids_are_preserved(self, temp_dir): ds = lance.dataset(uri) assert ds.count_rows() == 2 - base_paths = ds._ds.base_paths() + base_paths = ds._ds.base_paths() # type: ignore[attr-defined] assert 5 in base_paths assert 10 in base_paths - def test_mixed_explicit_and_implicit_ids(self, temp_dir): + def test_mixed_explicit_and_implicit_ids(self, temp_dir: str) -> None: """One base with explicit id, one without — no collision.""" base1_dir = Path(temp_dir) / "mixed_base1" base2_dir = Path(temp_dir) / "mixed_base2" @@ -628,14 +638,14 @@ def test_mixed_explicit_and_implicit_ids(self, temp_dir): ds = lance.dataset(uri) assert ds.count_rows() == 2 - base_paths = ds._ds.base_paths() + base_paths = ds._ds.base_paths() # type: ignore[attr-defined] assert 3 in base_paths base_ids = list(base_paths.keys()) assert len(set(base_ids)) == len(base_ids), ( f"Base path IDs must be unique, got: {base_paths}" ) - def test_dataset_root_base_gets_id_zero(self, temp_dir): + def test_dataset_root_base_gets_id_zero(self, temp_dir: str) -> None: """A base with is_dataset_root=True should receive id=0.""" base1_dir = Path(temp_dir) / "root_base" base2_dir = Path(temp_dir) / "extra_base" @@ -667,5 +677,5 @@ def test_dataset_root_base_gets_id_zero(self, temp_dir): ds = lance.dataset(uri) assert ds.count_rows() == 2 - base_paths = ds._ds.base_paths() + base_paths = ds._ds.base_paths() # type: ignore[attr-defined] assert 0 in base_paths diff --git a/tests/test_blob.py b/tests/test_blob.py index c36b2eb3..e2a6213c 100755 --- a/tests/test_blob.py +++ b/tests/test_blob.py @@ -15,7 +15,9 @@ import io import sys import tempfile +from collections.abc import Iterator from pathlib import Path +from typing import Any # Make local Lance python package available for import sys.path.insert( @@ -35,7 +37,7 @@ @pytest.fixture -def temp_dir(): +def temp_dir() -> Iterator[str]: with tempfile.TemporaryDirectory() as d: yield d @@ -68,7 +70,7 @@ def _generate_jpg_bytes() -> bytes: return os.urandom(64 * 64 * 3) -def test_single_blob_roundtrip(temp_dir): +def test_single_blob_roundtrip(temp_dir: str) -> None: """Test single blob column round-trip.""" path = Path(temp_dir) / "single_blob_roundtrip.lance" @@ -76,16 +78,15 @@ def test_single_blob_roundtrip(temp_dir): blob_values = [b"foo", b"bar", b"", None, b"\x00\x01\x02"] ids = pa.array([0, 1, 2, 3, 4], pa.int64()) - schema = pa.schema( - [ - pa.field( - "blob", - pa.large_binary(), - metadata={"lance-encoding:blob": "true"}, - ), - pa.field("id", pa.int64()), - ] - ) + schema_fields: list[pa.Field[Any]] = [ + pa.field( + "blob", + pa.large_binary(), + metadata={"lance-encoding:blob": "true"}, + ), + pa.field("id", pa.int64()), + ] + schema = pa.schema(schema_fields) table = pa.table( [pa.array(blob_values, type=pa.large_binary()), ids], schema=schema ) @@ -107,7 +108,7 @@ def test_single_blob_roundtrip(temp_dir): assert actual == expected -def test_multi_blob_roundtrip(temp_dir): +def test_multi_blob_roundtrip(temp_dir: str) -> None: """Test multiple blob columns round-trip.""" path = Path(temp_dir) / "multi_blob_roundtrip.lance" @@ -116,17 +117,12 @@ def test_multi_blob_roundtrip(temp_dir): blob2_values = [b"bar", b"baz", b"qux", None, b"\x02\x03\x04"] ids = pa.array([0, 1, 2, 3, 4], pa.int64()) - schema = pa.schema( - [ - pa.field( - "blob1", pa.large_binary(), metadata={"lance-encoding:blob": "true"} - ), - pa.field( - "blob2", pa.large_binary(), metadata={"lance-encoding:blob": "true"} - ), - pa.field("id", pa.int64()), - ] - ) + schema_fields: list[pa.Field[Any]] = [ + pa.field("blob1", pa.large_binary(), metadata={"lance-encoding:blob": "true"}), + pa.field("blob2", pa.large_binary(), metadata={"lance-encoding:blob": "true"}), + pa.field("id", pa.int64()), + ] + schema = pa.schema(schema_fields) table = pa.table( [ pa.array(blob1_values, type=pa.large_binary()), @@ -153,21 +149,20 @@ def test_multi_blob_roundtrip(temp_dir): @pytest.mark.lance_integration -def test_jpg_blob_integration(temp_dir, capsys): +def test_jpg_blob_integration( + temp_dir: str, capsys: pytest.CaptureFixture[str] +) -> None: """JPG blob integration test with real image data.""" # Prepare asset bytes jpg_bytes = _generate_jpg_bytes() # Build Arrow schema with blob + string + int - schema = pa.schema( - [ - pa.field( - "blob", pa.large_binary(), metadata={"lance-encoding:blob": "true"} - ), - pa.field("name", pa.string()), - pa.field("id", pa.int64()), - ] - ) + schema_fields: list[pa.Field[Any]] = [ + pa.field("blob", pa.large_binary(), metadata={"lance-encoding:blob": "true"}), + pa.field("name", pa.string()), + pa.field("id", pa.int64()), + ] + schema = pa.schema(schema_fields) # Construct a small table with repeated JPG bytes and varied name/id names = pa.array(["img_a", "img_b", "img_c"], type=pa.string()) @@ -207,23 +202,22 @@ def test_jpg_blob_integration(temp_dir, capsys): assert "jpg length=" in captured.out -def test_blob_projection_and_filter(temp_dir): +def test_blob_projection_and_filter(temp_dir: str) -> None: """Test projection and filtering on blob columns.""" path = Path(temp_dir) / "blob_projection_filter.lance" blob_values = [b"a", b"b", b"c", b"d"] ids = pa.array([10, 11, 12, 13], pa.int64()) - schema = pa.schema( - [ - pa.field( - "blob", - pa.large_binary(), - metadata={"lance-encoding:blob": "true"}, - ), - pa.field("id", pa.int64()), - ] - ) + schema_fields: list[pa.Field[Any]] = [ + pa.field( + "blob", + pa.large_binary(), + metadata={"lance-encoding:blob": "true"}, + ), + pa.field("id", pa.int64()), + ] + schema = pa.schema(schema_fields) table = pa.table( [pa.array(blob_values, type=pa.large_binary()), ids], schema=schema ) @@ -240,7 +234,7 @@ def test_blob_projection_and_filter(temp_dir): assert df["blob"].tolist() == [b"c", b"d"] -def test_multi_blob_projection_and_filter(temp_dir): +def test_multi_blob_projection_and_filter(temp_dir: str) -> None: """Test projection and filtering on multiple blob columns.""" path = Path(temp_dir) / "multi_blob_projection_filter.lance" @@ -249,17 +243,12 @@ def test_multi_blob_projection_and_filter(temp_dir): blob2_values = [b"bar", b"baz", b"qux", None, b"\x02\x03\x04"] ids = pa.array([0, 1, 2, 3, 4], pa.int64()) - schema = pa.schema( - [ - pa.field( - "blob1", pa.large_binary(), metadata={"lance-encoding:blob": "true"} - ), - pa.field( - "blob2", pa.large_binary(), metadata={"lance-encoding:blob": "true"} - ), - pa.field("id", pa.int64()), - ] - ) + schema_fields: list[pa.Field[Any]] = [ + pa.field("blob1", pa.large_binary(), metadata={"lance-encoding:blob": "true"}), + pa.field("blob2", pa.large_binary(), metadata={"lance-encoding:blob": "true"}), + pa.field("id", pa.int64()), + ] + schema = pa.schema(schema_fields) table = pa.table( [ pa.array(blob1_values, type=pa.large_binary()), @@ -284,7 +273,7 @@ def test_multi_blob_projection_and_filter(temp_dir): assert df["blob2"].tolist() == expected_blob2 -def test_stream_copy_basic_local(temp_dir): +def test_stream_copy_basic_local(temp_dir: str) -> None: """Basic streaming copy on local filesystem (batch_size=1) with legacy blob data.""" import lance @@ -292,18 +281,17 @@ def test_stream_copy_basic_local(temp_dir): dst_path = Path(temp_dir) / "dst_small_legacy_blob_copy.lance" # Build Arrow schema with a legacy blob column + some regular columns - schema = pa.schema( - [ - pa.field( - "blob", - pa.large_binary(), - metadata={"lance-encoding:blob": "true"}, - ), - pa.field("id", pa.int64()), - pa.field("name", pa.string()), - pa.field("val", pa.float64()), - ] - ) + schema_fields: list[pa.Field[Any]] = [ + pa.field( + "blob", + pa.large_binary(), + metadata={"lance-encoding:blob": "true"}, + ), + pa.field("id", pa.int64()), + pa.field("name", pa.string()), + pa.field("val", pa.float64()), + ] + schema = pa.schema(schema_fields) blob_values = [b"x", None, b"", b"y"] ids = pa.array([1, 2, 3, 4], pa.int64()) @@ -344,23 +332,20 @@ def test_stream_copy_basic_local(temp_dir): pd.testing.assert_frame_equal(src_df, dst_df) -def test_stream_copy_resume_local(temp_dir): +def test_stream_copy_resume_local(temp_dir: str) -> None: """Resume streaming copy with legacy blob data: write first 2 rows then append the rest.""" src_path = Path(temp_dir) / "src_resume_legacy_blob.lance" dst_path = Path(temp_dir) / "dst_resume_legacy_blob_copy.lance" # Legacy blob schema - schema = pa.schema( - [ - pa.field( - "blob", pa.large_binary(), metadata={"lance-encoding:blob": "true"} - ), - pa.field("id", pa.int64()), - pa.field("name", pa.string()), - pa.field("val", pa.float64()), - ] - ) + schema_fields: list[pa.Field[Any]] = [ + pa.field("blob", pa.large_binary(), metadata={"lance-encoding:blob": "true"}), + pa.field("id", pa.int64()), + pa.field("name", pa.string()), + pa.field("val", pa.float64()), + ] + schema = pa.schema(schema_fields) blob_values = [b"m", b"n", b"o", None, b"p"] ids = pa.array([1, 2, 3, 4, 5], pa.int64()) diff --git a/tests/test_blob_v2.py b/tests/test_blob_v2.py index 60ac7436..a1b7a01b 100755 --- a/tests/test_blob_v2.py +++ b/tests/test_blob_v2.py @@ -18,6 +18,7 @@ import inspect import sys from pathlib import Path +from typing import Any import pyarrow as pa import pytest @@ -72,7 +73,7 @@ def _build_blob_v2_table( ) sig = inspect.signature(Blob.from_uri) - blob_kwargs: dict[str, object] = {} + blob_kwargs: dict[str, Any] = {} if "position" in sig.parameters and "size" in sig.parameters: blob_kwargs = {"position": 0, "size": len(external_payload)} @@ -122,7 +123,7 @@ def _build_external_only_blob_v2_table( ) sig = inspect.signature(Blob.from_uri) - blob_kwargs: dict[str, object] = {} + blob_kwargs: dict[str, Any] = {} if "position" in sig.parameters and "size" in sig.parameters: blob_kwargs = {"position": 0, "size": len(payload)} @@ -288,7 +289,7 @@ def test_blob_v2_external_blob_ingest_write_lance( ) path = tmp_path / f"blob_v2_external_ingest_{stream}.lance" - stream_kwargs = {"stream": stream} + stream_kwargs: dict[str, Any] = {"stream": stream} if stream: stream_kwargs["batch_size"] = 1 @@ -358,7 +359,7 @@ def test_blob_v2_reference_multi_base_all_lance_ray_paths(tmp_path: Path) -> Non table, inline_payload, external_payload, external_base, _ = _build_blob_v2_table( tmp_path ) - base_store_params = {external_base.as_uri(): {}} + base_store_params: dict[str, dict[str, Any]] = {external_base.as_uri(): {}} initial_bases = _initial_bases(external_base) # Non-streaming write_lance + full/projection/filter/fragment reads. @@ -510,7 +511,10 @@ def test_blob_v2_create_with_initial_bases_and_target_bases(tmp_path: Path) -> N _build_multi_base_blob_v2_table(tmp_path) ) initial_bases = _multi_initial_bases(base_a, base_b) - base_store_params = {base_a.as_uri(): {}, base_b.as_uri(): {}} + base_store_params: dict[str, dict[str, Any]] = { + base_a.as_uri(): {}, + base_b.as_uri(): {}, + } path = tmp_path / "blob_v2_create_target_bases.lance" @@ -572,7 +576,10 @@ def test_blob_v2_target_bases_multi_base_routing(tmp_path: Path) -> None: _build_multi_base_blob_v2_table(tmp_path) ) initial_bases = _multi_initial_bases(base_a, base_b) - base_store_params = {base_a.as_uri(): {}, base_b.as_uri(): {}} + base_store_params: dict[str, dict[str, Any]] = { + base_a.as_uri(): {}, + base_b.as_uri(): {}, + } path = tmp_path / "blob_v2_multi_base_routing.lance" @@ -679,7 +686,10 @@ def test_blob_v2_append_with_target_bases_stream(tmp_path: Path) -> None: _build_multi_base_blob_v2_table(tmp_path) ) initial_bases = _multi_initial_bases(base_a, base_b) - base_store_params = {base_a.as_uri(): {}, base_b.as_uri(): {}} + base_store_params: dict[str, dict[str, Any]] = { + base_a.as_uri(): {}, + base_b.as_uri(): {}, + } path = tmp_path / "blob_v2_target_bases_stream.lance" diff --git a/tests/test_fragment.py b/tests/test_fragment.py index bf9dded8..ad423e5f 100644 --- a/tests/test_fragment.py +++ b/tests/test_fragment.py @@ -2,6 +2,7 @@ import warnings from pathlib import Path +from typing import Any, Optional, cast import lance import lance_ray.io as lr @@ -12,21 +13,25 @@ from lance_ray.fragment import LanceFragmentWriter -def _legacy_write_fragments(reader, uri, *, schema=None): +def _legacy_write_fragments( + reader: Any, uri: Any, *, schema: Optional[pa.Schema] = None +) -> list[Any]: return [] def _write_fragments_with_external_blob_options( - reader, - uri, + reader: Any, + uri: Any, *, - external_blob_mode="reference", - allow_external_blob_outside_bases=False, -): + external_blob_mode: str = "reference", + allow_external_blob_outside_bases: bool = False, +) -> list[Any]: return [] -def test_fragment_writer_external_blob_options_fail_fast(monkeypatch, tmp_path: Path): +def test_fragment_writer_external_blob_options_fail_fast( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: import lance.fragment as lance_fragment monkeypatch.setattr( @@ -53,7 +58,9 @@ def test_fragment_writer_external_blob_options_fail_fast(monkeypatch, tmp_path: ) -def test_datasink_external_blob_options_fail_fast(monkeypatch, tmp_path: Path): +def test_datasink_external_blob_options_fail_fast( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: import lance.fragment as lance_fragment monkeypatch.setattr( @@ -66,7 +73,9 @@ def test_datasink_external_blob_options_fail_fast(monkeypatch, tmp_path: Path): LanceDatasink(str(tmp_path), external_blob_mode="ingest") -def test_write_lance_external_blob_options_fail_fast(monkeypatch, tmp_path: Path): +def test_write_lance_external_blob_options_fail_fast( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: import lance.fragment as lance_fragment monkeypatch.setattr( @@ -76,13 +85,13 @@ def test_write_lance_external_blob_options_fail_fast(monkeypatch, tmp_path: Path ) with pytest.raises(RuntimeError, match="external_blob_mode.*write_fragments"): - lr.write_lance(object(), str(tmp_path), external_blob_mode="ingest") + lr.write_lance(cast(Any, object()), str(tmp_path), external_blob_mode="ingest") def test_base_store_params_fail_fast_when_fragment_api_unsupported( - monkeypatch, + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, -): +) -> None: import lance.fragment as lance_fragment monkeypatch.setattr( @@ -90,7 +99,7 @@ def test_base_store_params_fail_fast_when_fragment_api_unsupported( "write_fragments", _legacy_write_fragments, ) - base_store_params = {tmp_path.as_uri(): {}} + base_store_params: dict[str, dict[str, Any]] = {tmp_path.as_uri(): {}} with pytest.raises(RuntimeError, match="base_store_params.*write_fragments"): LanceFragmentWriter( @@ -103,13 +112,15 @@ def test_base_store_params_fail_fast_when_fragment_api_unsupported( LanceDatasink(str(tmp_path), base_store_params=base_store_params) with pytest.raises(RuntimeError, match="base_store_params.*write_fragments"): - lr.write_lance(object(), str(tmp_path), base_store_params=base_store_params) + lr.write_lance( + cast(Any, object()), str(tmp_path), base_store_params=base_store_params + ) def test_target_bases_fail_fast_when_fragment_api_unsupported( - monkeypatch, + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, -): +) -> None: import lance.fragment as lance_fragment monkeypatch.setattr( @@ -130,13 +141,13 @@ def test_target_bases_fail_fast_when_fragment_api_unsupported( LanceDatasink(str(tmp_path), target_bases=target_bases) with pytest.raises(RuntimeError, match="target_bases.*write_fragments"): - lr.write_lance(object(), str(tmp_path), target_bases=target_bases) + lr.write_lance(cast(Any, object()), str(tmp_path), target_bases=target_bases) def test_allow_external_blob_outside_bases_ignored_for_ingest( - monkeypatch, + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, -): +) -> None: import lance.fragment as lance_fragment monkeypatch.setattr( @@ -157,9 +168,9 @@ def test_allow_external_blob_outside_bases_ignored_for_ingest( def test_unsupported_ingest_with_allow_external_blob_outside_bases_does_not_warn( - monkeypatch, + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, -): +) -> None: import lance.fragment as lance_fragment monkeypatch.setattr( @@ -185,16 +196,22 @@ class TestLanceFragmentWriterCommitter: """Test cases for LanceFragmentWriter and LanceCommitter.""" @pytest.mark.filterwarnings("ignore::DeprecationWarning") - def test_fragment_writer_committer(self, tmp_path: Path): + def test_fragment_writer_committer(self, tmp_path: Path) -> None: """Test fragment writer and committer for large-scale data.""" - schema = pa.schema([pa.field("id", pa.int64()), pa.field("str", pa.string())]) + schema_fields: list[pa.Field[Any]] = [ + pa.field("id", pa.int64()), + pa.field("str", pa.string()), + ] + schema = pa.schema(schema_fields) # Use fragment writer and committer ( ray.data.range(10) .map(lambda x: {"id": x["id"], "str": f"str-{x['id']}"}) - .map_batches(LanceFragmentWriter(tmp_path, schema=schema), batch_size=5) - .write_datasink(LanceFragmentCommitter(tmp_path)) + .map_batches( + LanceFragmentWriter(str(tmp_path), schema=schema), batch_size=5 + ) + .write_datasink(LanceFragmentCommitter(str(tmp_path))) ) # Verify the dataset @@ -203,20 +220,23 @@ def test_fragment_writer_committer(self, tmp_path: Path): assert ds.schema == schema tbl = ds.to_table() - assert sorted(tbl["id"].to_pylist()) == list(range(10)) + assert sorted(cast("list[int]", tbl["id"].to_pylist())) == list(range(10)) assert set(tbl["str"].to_pylist()) == set([f"str-{i}" for i in range(10)]) # Should have 2 fragments since batch_size=5 and we have 10 rows assert len(ds.get_fragments()) == 2 @pytest.mark.filterwarnings("ignore::DeprecationWarning") - def test_fragment_writer_committer_enables_stable_row_ids(self, tmp_path: Path): - schema = pa.schema([pa.field("id", pa.int64())]) + def test_fragment_writer_committer_enables_stable_row_ids( + self, tmp_path: Path + ) -> None: + schema_fields: list[pa.Field[Any]] = [pa.field("id", pa.int64())] + schema = pa.schema(schema_fields) ( ray.data.range(10) .map_batches( LanceFragmentWriter( - tmp_path, + str(tmp_path), schema=schema, enable_stable_row_ids=True, ), @@ -224,7 +244,7 @@ def test_fragment_writer_committer_enables_stable_row_ids(self, tmp_path: Path): ) .write_datasink( LanceFragmentCommitter( - tmp_path, + str(tmp_path), enable_stable_row_ids=True, ) ) @@ -255,15 +275,14 @@ def test_fragment_writer_committer_enables_stable_row_ids(self, tmp_path: Path): assert after == before @pytest.mark.filterwarnings("ignore::DeprecationWarning") - def test_fragment_writer_with_transform(self, tmp_path: Path): + def test_fragment_writer_with_transform(self, tmp_path: Path) -> None: """Test fragment writer with custom transform function.""" - schema = pa.schema( - [ - pa.field("id", pa.int64()), - pa.field("str", pa.string()), - pa.field("doubled", pa.int64()), - ] - ) + schema_fields: list[pa.Field[Any]] = [ + pa.field("id", pa.int64()), + pa.field("str", pa.string()), + pa.field("doubled", pa.int64()), + ] + schema = pa.schema(schema_fields) def transform(batch: pa.Table) -> pa.Table: """Transform function to add a doubled column.""" @@ -276,10 +295,10 @@ def transform(batch: pa.Table) -> pa.Table: ray.data.range(5) .map(lambda x: {"id": x["id"], "str": f"str-{x['id']}"}) .map_batches( - LanceFragmentWriter(tmp_path, schema=schema, transform=transform), + LanceFragmentWriter(str(tmp_path), schema=schema, transform=transform), batch_size=5, ) - .write_datasink(LanceFragmentCommitter(tmp_path)) + .write_datasink(LanceFragmentCommitter(str(tmp_path))) ) # Verify the dataset @@ -291,16 +310,20 @@ def transform(batch: pa.Table) -> pa.Table: assert tbl_sorted.column("doubled").to_pylist() == [0, 2, 4, 6, 8] @pytest.mark.filterwarnings("ignore::DeprecationWarning") - def test_fragment_writer_append_mode(self, tmp_path: Path): + def test_fragment_writer_append_mode(self, tmp_path: Path) -> None: """Test fragment writer with append mode.""" - schema = pa.schema([pa.field("id", pa.int64()), pa.field("str", pa.string())]) + schema_fields: list[pa.Field[Any]] = [ + pa.field("id", pa.int64()), + pa.field("str", pa.string()), + ] + schema = pa.schema(schema_fields) # Write initial data ( ray.data.range(5) .map(lambda x: {"id": x["id"], "str": f"str-{x['id']}"}) - .map_batches(LanceFragmentWriter(tmp_path, schema=schema)) - .write_datasink(LanceFragmentCommitter(tmp_path, mode="create")) + .map_batches(LanceFragmentWriter(str(tmp_path), schema=schema)) + .write_datasink(LanceFragmentCommitter(str(tmp_path), mode="create")) ) # Append more data @@ -308,28 +331,32 @@ def test_fragment_writer_append_mode(self, tmp_path: Path): ray.data.range(10) .filter(lambda row: row["id"] >= 5) .map(lambda x: {"id": x["id"], "str": f"str-{x['id']}"}) - .map_batches(LanceFragmentWriter(tmp_path, schema=schema)) - .write_datasink(LanceFragmentCommitter(tmp_path, mode="append")) + .map_batches(LanceFragmentWriter(str(tmp_path), schema=schema)) + .write_datasink(LanceFragmentCommitter(str(tmp_path), mode="append")) ) # Verify the dataset ds = lance.dataset(tmp_path) assert ds.count_rows() == 10 tbl = ds.to_table() - assert sorted(tbl["id"].to_pylist()) == list(range(10)) + assert sorted(cast("list[int]", tbl["id"].to_pylist())) == list(range(10)) @pytest.mark.filterwarnings("ignore::DeprecationWarning") - def test_fragment_writer_empty_write(self, tmp_path: Path): + def test_fragment_writer_empty_write(self, tmp_path: Path) -> None: """Test fragment writer with empty data.""" - schema = pa.schema([pa.field("id", pa.int64()), pa.field("str", pa.string())]) + schema_fields: list[pa.Field[Any]] = [ + pa.field("id", pa.int64()), + pa.field("str", pa.string()), + ] + schema = pa.schema(schema_fields) # Write empty data (filter everything out) ( ray.data.range(10) .filter(lambda row: row["id"] > 10) # Filter out everything .map(lambda x: {"id": x["id"], "str": f"str-{x['id']}"}) - .map_batches(LanceFragmentWriter(tmp_path, schema=schema)) - .write_datasink(LanceFragmentCommitter(tmp_path)) + .map_batches(LanceFragmentWriter(str(tmp_path), schema=schema)) + .write_datasink(LanceFragmentCommitter(str(tmp_path))) ) # Empty write should not create a dataset @@ -337,22 +364,26 @@ def test_fragment_writer_empty_write(self, tmp_path: Path): lance.dataset(tmp_path) @pytest.mark.filterwarnings("ignore::DeprecationWarning") - def test_fragment_writer_none_values(self, tmp_path: Path): + def test_fragment_writer_none_values(self, tmp_path: Path) -> None: """Test fragment writer with None values.""" - def create_row(row): + def create_row(row: dict[str, Any]) -> dict[str, Any]: return { "id": row["id"], "str": None if row["id"] % 2 == 0 else f"str-{row['id']}", } - schema = pa.schema([pa.field("id", pa.int64()), pa.field("str", pa.string())]) + schema_fields: list[pa.Field[Any]] = [ + pa.field("id", pa.int64()), + pa.field("str", pa.string()), + ] + schema = pa.schema(schema_fields) ( ray.data.range(10) .map(create_row) - .map_batches(LanceFragmentWriter(tmp_path, schema=schema)) - .write_datasink(LanceFragmentCommitter(tmp_path)) + .map_batches(LanceFragmentWriter(str(tmp_path), schema=schema)) + .write_datasink(LanceFragmentCommitter(str(tmp_path))) ) # Verify the dataset @@ -362,7 +393,9 @@ def create_row(row): str_values = tbl["str"].to_pylist() id_values = tbl["id"].to_pylist() # Even IDs should have None values - for id_val, str_val in zip(id_values, str_values, strict=False): + for id_val, str_val in zip( + cast("list[int]", id_values), str_values, strict=False + ): if id_val % 2 == 0: # None values might be represented as None or as 'nan' string assert str_val is None or str(str_val) == "nan", (