From eb0a770bf4ff936584434f660c36d3b8ddb6766f Mon Sep 17 00:00:00 2001 From: fecet Date: Wed, 18 Dec 2024 20:31:12 +0800 Subject: [PATCH] add test --- python/python/tests/test_balanced.py | 24 ++++++++++++++++++++++++ python/python/tests/test_fragment.py | 4 ++-- 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/python/python/tests/test_balanced.py b/python/python/tests/test_balanced.py index a7d33bd3d0..9259a162a0 100644 --- a/python/python/tests/test_balanced.py +++ b/python/python/tests/test_balanced.py @@ -58,6 +58,30 @@ def balanced_dataset(tmp_path, big_val): ) +def test_write_fragments(balanced_dataset, tmp_path): + tbl = balanced_dataset._take_rows(range(10)) + default_frags, blob_frags = lance.fragment.write_fragments( + tbl, + tmp_path / "ds", + with_blobs=True, + enable_move_stable_row_ids=True, + ) + operation = lance.LanceOperation.Overwrite(tbl.schema, default_frags) + blob_operation = lance.LanceOperation.Overwrite( + pa.schema([tbl.schema.field("blobs")]), blob_frags + ) + + ds = lance.LanceDataset.commit( + tmp_path / "ds", + operation, + blobs_operation=blob_operation, + enable_v2_manifest_paths=True, + ) + dataset = lance.LanceDataset(tmp_path / "ds") + + assert dataset._take_rows(range(10)) == balanced_dataset._take_rows(range(10)) + + def test_append_then_take(balanced_dataset, tmp_path, big_val): blob_dir = tmp_path / "test_ds" / "_blobs" / "data" assert len(list(blob_dir.iterdir())) == 8 diff --git a/python/python/tests/test_fragment.py b/python/python/tests/test_fragment.py index c14cba5a73..5a18c3de77 100644 --- a/python/python/tests/test_fragment.py +++ b/python/python/tests/test_fragment.py @@ -317,7 +317,7 @@ def test_create_from_file(tmp_path): frag = LanceFragment.create_from_file(fragment_name, dataset, 0) op = LanceOperation.Append([frag]) - dataset = lance.LanceDataset.commit(dataset.uri, op, dataset.version) + dataset = lance.LanceDataset.commit(dataset.uri, op, read_version=dataset.version) frag = dataset.get_fragments()[0] assert frag.fragment_id == 0 @@ -331,7 +331,7 @@ def test_create_from_file(tmp_path): frag = LanceFragment.create_from_file(fragment_name, dataset, 0) op = LanceOperation.Append([frag]) - dataset = lance.LanceDataset.commit(dataset.uri, op, dataset.version) + dataset = lance.LanceDataset.commit(dataset.uri, op, read_version=dataset.version) frag = dataset.get_fragments()[1] assert frag.fragment_id == 1