diff --git a/python/python/ci_benchmarks/benchmarks/test_overlay_manifest.py b/python/python/ci_benchmarks/benchmarks/test_overlay_manifest.py new file mode 100644 index 00000000000..0303e41c7a3 --- /dev/null +++ b/python/python/ci_benchmarks/benchmarks/test_overlay_manifest.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The Lance Authors + +"""Benchmark #1: impact of overlay files on manifest size. + +Each committed overlay adds a ``DataOverlayFile`` entry (a value-file pointer +plus a serialized coverage bitmap) to every overlaid fragment's metadata in the +manifest. This measures how the manifest grows with the number of overlays and +the coverage size, since manifests are read on every dataset open. +""" + +import pytest +from ci_benchmarks.overlays import ( + commit_overlay_layers, + make_base_dataset, + manifest_size, +) + +NUM_ROWS = 1_000_000 +# Single fragment so every overlay lands on the same fragment's metadata. +ROWS_PER_FILE = NUM_ROWS + + +@pytest.mark.parametrize("num_overlays", [0, 1, 4, 16, 64]) +@pytest.mark.parametrize( + "fraction,pattern", + [(0.01, "contiguous"), (0.01, "stride"), (0.1, "stride")], + ids=["1pct-contiguous", "1pct-stride", "10pct-stride"], +) +def test_overlay_manifest_size( + tmp_path, record_property, num_overlays, fraction, pattern +): + base = str(tmp_path / "ds") + ds = make_base_dataset(base, NUM_ROWS, ROWS_PER_FILE, "int32", "2.1") + base_bytes = manifest_size(ds) + + if num_overlays: + ds = commit_overlay_layers(ds, num_overlays, fraction, pattern, "int32") + + total = manifest_size(ds) + growth = total - base_bytes + # Guard the fixture: committed overlays must enlarge the manifest, else the + # benchmark would report growth=0 for overlays that were never recorded. + if num_overlays: + assert growth > 0, "overlays did not grow the manifest" + per_overlay = growth / num_overlays if num_overlays else 0 + + record_property("manifest_bytes", total) + record_property("manifest_growth_bytes", growth) + record_property("bytes_per_overlay", per_overlay) diff --git a/python/python/ci_benchmarks/benchmarks/test_overlay_read.py b/python/python/ci_benchmarks/benchmarks/test_overlay_read.py new file mode 100644 index 00000000000..27038d105ce --- /dev/null +++ b/python/python/ci_benchmarks/benchmarks/test_overlay_read.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The Lance Authors + +"""Benchmark #3: impact of overlay files on take and scan workloads. + +On read, every overlay covering a requested cell must be consulted and its value +merged over the base. This measures how take and full-scan cost scale with: +- the number of overlay layers stacked on a fragment (compaction payoff), +- coverage fraction (how many cells are overlaid), +- fragmentation (contiguous run vs. strided -> pages touched), and +- value width (a 4-byte int32 vs. a wide fixed-size-list embedding). + +Wall time is measured warm via pytest-benchmark; read IO (bytes + IOPS) is +measured once cold, after dropping the page cache, via ``io_stats_incremental``. +""" + +import random + +import lance +import pytest +from ci_benchmarks.overlays import commit_overlay_layers, make_base_dataset +from ci_benchmarks.utils import wipe_os_cache + +NUM_ROWS = 1_000_000 +ROWS_PER_FILE = NUM_ROWS # single fragment: isolates overlay-layer scaling +TAKE_ROWS = 100 + +# Wide value column: a 3072-d float32 embedding is 12 KiB/row, ~750x an int32 +# cell. Fewer rows keep the base file to ~1.2 GiB while each cell still dominates +# read cost, so the merge/interleave a scan pays per overlay layer moves real +# payload rather than 4-byte integers. +WIDE_EMBEDDING_DIM = 3072 +NUM_ROWS_WIDE = 100_000 + + +def _take_indices(num_rows: int) -> list: + rng = random.Random(0) + return sorted(rng.sample(range(num_rows), TAKE_ROWS)) + + +def _covered_value(ds: lance.LanceDataset): + """``val`` at offset 0, which every coverage pattern includes. + + Used to guard the fixture: after committing overlays a covered cell must + read back a new value, otherwise a read-invisible overlay would let the + benchmark silently time plain base reads. + """ + return ds.take([0], columns=["val"]).column("val").to_pylist()[0] + + +def _measure_cold_io(ds: lance.LanceDataset, base: str, work): + """Drop the page cache, run ``work`` once, return its read IO stats.""" + wipe_os_cache(base) + ds.io_stats_incremental() # reset + work() + stats = ds.io_stats_incremental() + return stats.read_bytes, stats.read_iops + + +def _run_read(benchmark, record_property, base, ds, workload, num_rows): + if workload == "take": + indices = _take_indices(num_rows) + + def work(): + ds.take(indices, columns=["val"]) + else: + + def work(): + ds.to_table(columns=["val"]) + + read_bytes, read_iops = _measure_cold_io(ds, base, work) + record_property("cold_read_bytes", read_bytes) + record_property("cold_read_iops", read_iops) + + benchmark(work) + + +@pytest.mark.parametrize("version", ["2.0", "2.1"]) +@pytest.mark.parametrize("workload", ["take", "scan"]) +@pytest.mark.parametrize("num_overlays", [0, 4, 16]) +@pytest.mark.parametrize( + "fraction,pattern", + [(0.01, "contiguous"), (0.01, "stride")], + ids=["1pct-contiguous", "1pct-stride"], +) +def test_overlay_read_scaling( + benchmark, + tmp_path, + record_property, + version, + workload, + num_overlays, + fraction, + pattern, +): + base = str(tmp_path / "ds") + ds = make_base_dataset(base, NUM_ROWS, ROWS_PER_FILE, "int32", version) + if num_overlays: + base_val = _covered_value(ds) + ds = commit_overlay_layers(ds, num_overlays, fraction, pattern, "int32") + assert _covered_value(ds) != base_val, "overlay not visible on read" + _run_read(benchmark, record_property, base, ds, workload, NUM_ROWS) + + +# Mirror test_overlay_read_scaling but on a wide 3072-d embedding column, so the +# take/scan-vs-layers story can be read for a fat value column rather than a +# 4-byte one. Pinned to the 2.1 format (v2.0 layer scaling is already covered by +# the narrow scan above; the wide-column question is about per-layer payload). +@pytest.mark.parametrize("workload", ["take", "scan"]) +@pytest.mark.parametrize("num_overlays", [0, 4, 16]) +@pytest.mark.parametrize( + "fraction,pattern", + [(0.01, "contiguous"), (0.01, "stride")], + ids=["1pct-contiguous", "1pct-stride"], +) +def test_overlay_read_wide( + benchmark, + tmp_path, + record_property, + workload, + num_overlays, + fraction, + pattern, +): + base = str(tmp_path / "ds") + ds = make_base_dataset( + base, + NUM_ROWS_WIDE, + NUM_ROWS_WIDE, + "embedding", + "2.1", + embedding_dim=WIDE_EMBEDDING_DIM, + ) + if num_overlays: + base_val = _covered_value(ds) + ds = commit_overlay_layers( + ds, + num_overlays, + fraction, + pattern, + "embedding", + embedding_dim=WIDE_EMBEDDING_DIM, + ) + assert _covered_value(ds) != base_val, "overlay not visible on read" + _run_read(benchmark, record_property, base, ds, workload, NUM_ROWS_WIDE) diff --git a/python/python/ci_benchmarks/overlays.py b/python/python/ci_benchmarks/overlays.py new file mode 100644 index 00000000000..c2610e36d2e --- /dev/null +++ b/python/python/ci_benchmarks/overlays.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The Lance Authors + +"""Shared helpers for the data-overlay benchmark suite. + +Data overlay files supply replacement values for a subset of (row offset, field) +cells in a fragment, merged on read, without rewriting the base data file. These +helpers build synthetic base datasets, commit overlay layers through the public +``lance.LanceOperation.DataOverlay`` operation, and measure their cost. +""" + +import os + +# Data overlay support is gated off in release builds unless this is set (it is +# always on in debug builds). Benchmarks usually run against a release build, so +# enable it here, before lance is imported, or reading overlay datasets fails. +os.environ.setdefault("LANCE_ENABLE_UNSTABLE_DATA_OVERLAY_FILES", "1") + +from typing import List # noqa: E402 +from urllib.parse import urlparse # noqa: E402 + +import lance # noqa: E402 +import numpy as np # noqa: E402 +import pyarrow as pa # noqa: E402 +from lance._datagen import rand_batches # noqa: E402 +from lance.file import LanceFileWriter # noqa: E402 + +# Default width for the ``embedding`` dtype. Individual benchmarks override it +# via ``embedding_dim`` to model narrow vs. wide (e.g. 3072-d) value columns. +EMBEDDING_DIM = 128 + + +def _value_type(dtype: str, embedding_dim: int = EMBEDDING_DIM) -> pa.DataType: + if dtype == "int32": + return pa.int32() + if dtype == "embedding": + return pa.list_(pa.float32(), embedding_dim) + raise ValueError(f"unknown overlay benchmark dtype {dtype!r}") + + +def _gen_values( + dtype: str, + n: int, + rng: np.random.Generator, + embedding_dim: int = EMBEDDING_DIM, +) -> pa.Array: + if dtype == "int32": + return pa.array(rng.integers(0, 1 << 30, size=n, dtype=np.int32)) + if dtype == "embedding": + flat = rng.random(n * embedding_dim, dtype=np.float32) + return pa.FixedSizeListArray.from_arrays(pa.array(flat), embedding_dim) + raise ValueError(f"unknown overlay benchmark dtype {dtype!r}") + + +def make_base_dataset( + base_path: str, + num_rows: int, + rows_per_file: int, + dtype: str, + version: str, + embedding_dim: int = EMBEDDING_DIM, +) -> lance.LanceDataset: + """Create a base dataset with an ``id`` column and a ``val`` column. + + ``val`` is the column overlays target. ``rows_per_file`` controls the number + of fragments (``num_rows // rows_per_file``), which must divide ``num_rows``. + ``embedding_dim`` sets the width of the ``val`` column when ``dtype`` is + ``embedding`` (e.g. 3072 for a wide embedding). + """ + if num_rows % rows_per_file: + raise ValueError( + f"num_rows ({num_rows}) must be a multiple of rows_per_file " + f"({rows_per_file})" + ) + schema = pa.schema({"id": pa.int64(), "val": _value_type(dtype, embedding_dim)}) + # One fragment per batch; lance-datagen fills both columns with random data. + reader = rand_batches( + schema, num_batches=num_rows // rows_per_file, rows_per_batch=rows_per_file + ) + return lance.write_dataset( + reader, + base_path, + max_rows_per_file=rows_per_file, + data_storage_version=version, + ) + + +def coverage_offsets(num_rows: int, fraction: float, pattern: str) -> List[int]: + """Offsets within a fragment covered by an overlay. + + ``contiguous`` packs the covered cells into a single leading run (few pages + touched); ``stride`` spreads them evenly across the fragment (many pages + touched). Both cover ``round(num_rows * fraction)`` cells. + """ + count = max(1, int(round(num_rows * fraction))) + if pattern == "contiguous": + return list(range(count)) + if pattern == "stride": + step = max(1, num_rows // count) + return list(range(0, num_rows, step))[:count] + raise ValueError(f"unknown coverage pattern {pattern!r}") + + +def _val_field_id(ds: lance.LanceDataset) -> int: + base_df = ds.get_fragments()[0].metadata.files[0] + names = [f.name for f in ds.schema] + return base_df.fields[names.index("val")] + + +def commit_overlay_layers( + ds: lance.LanceDataset, + num_layers: int, + fraction: float, + pattern: str, + dtype: str, + *, + seed: int = 0, + embedding_dim: int = EMBEDDING_DIM, +) -> lance.LanceDataset: + """Commit ``num_layers`` overlays on ``val``, each covering the same offsets + in every fragment so that all layers must be consulted on read (the case + that motivates compaction). Returns the updated dataset. + """ + base_df = ds.get_fragments()[0].metadata.files[0] + field_id = _val_field_id(ds) + data_dir = os.path.join(_local_path(ds), "data") + for layer in range(num_layers): + rng = np.random.default_rng(seed + layer + 1) + groups = [] + for frag in ds.get_fragments(): + offsets = coverage_offsets(frag.count_rows(), fraction, pattern) + values = _gen_values(dtype, len(offsets), rng, embedding_dim) + batch = pa.record_batch([values], names=["val"]) + name = f"overlay_l{layer}_f{frag.fragment_id}.lance" + path = os.path.join(data_dir, name) + with LanceFileWriter(path) as writer: + writer.write_batch(batch) + df = lance.fragment.DataFile( + path=name, + fields=[field_id], + column_indices=[0], + file_major_version=base_df.file_major_version, + file_minor_version=base_df.file_minor_version, + file_size_bytes=os.path.getsize(path), + ) + groups.append( + lance.LanceOperation.DataOverlayGroup( + fragment_id=frag.fragment_id, + overlays=[ + lance.LanceOperation.DataOverlayFile( + data_file=df, offsets=offsets + ) + ], + ) + ) + op = lance.LanceOperation.DataOverlay(groups=groups) + ds = lance.LanceDataset.commit(ds, op, read_version=ds.version) + return ds + + +# --- Measurement helpers ---------------------------------------------------- + + +def _local_path(ds: lance.LanceDataset) -> str: + parsed = urlparse(ds.uri) + return parsed.path if parsed.scheme == "file" else ds.uri + + +def manifest_size(ds: lance.LanceDataset) -> int: + """Size in bytes of the manifest for the dataset's current version.""" + # Manifests are named `{u64::MAX - version}.manifest` so that a plain + # lexicographic directory listing yields newest-version-first. + name = f"{(1 << 64) - 1 - ds.version}.manifest" + return os.path.getsize(os.path.join(_local_path(ds), "_versions", name)) diff --git a/python/python/lance/_datagen.py b/python/python/lance/_datagen.py index b156066eca6..45826503c5b 100644 --- a/python/python/lance/_datagen.py +++ b/python/python/lance/_datagen.py @@ -21,10 +21,13 @@ def rand_batches( *, num_batches: Optional[int] = None, batch_size_bytes: Optional[int] = None, + rows_per_batch: Optional[int] = None, ): if not datagen.is_datagen_supported(): raise NotImplementedError( "This version of lance was not built with the datagen feature" ) - batch_iter = datagen.rand_batches(schema, num_batches, batch_size_bytes) + batch_iter = datagen.rand_batches( + schema, num_batches, batch_size_bytes, rows_per_batch + ) return pa.RecordBatchReader.from_batches(schema, batch_iter) diff --git a/python/src/datagen.rs b/python/src/datagen.rs index 8b046c37f24..194ba3d6a73 100644 --- a/python/src/datagen.rs +++ b/python/src/datagen.rs @@ -1,7 +1,7 @@ use arrow::pyarrow::PyArrowType; -use arrow_array::RecordBatch; +use arrow_array::{RecordBatch, RecordBatchReader}; use arrow_schema::Schema; -use lance_datagen::{BatchCount, ByteCount}; +use lance_datagen::{BatchCount, ByteCount, RowCount}; use pyo3::{ Bound, PyResult, Python, pyfunction, types::{PyModule, PyModuleMethods}, @@ -16,22 +16,44 @@ pub fn is_datagen_supported() -> bool { true } +/// Generate `batch_count` batches of random data for `schema`. +/// +/// Batch size is set either by `rows_in_batch` (exact rows per batch) or +/// `bytes_in_batch` (approximate bytes per batch); the two are mutually +/// exclusive. When neither is given the byte-based default is used. #[pyfunction] -#[pyo3(signature=(schema, batch_count=None, bytes_in_batch=None))] +#[pyo3(signature=(schema, batch_count=None, bytes_in_batch=None, rows_in_batch=None))] pub fn rand_batches( schema: PyArrowType, batch_count: Option, bytes_in_batch: Option, + rows_in_batch: Option, ) -> PyResult>> { - lance_datagen::rand(&schema.0) - .into_reader_bytes( - ByteCount::from(bytes_in_batch.unwrap_or(DEFAULT_BATCH_SIZE_BYTES)), - BatchCount::from(batch_count.unwrap_or(DEFAULT_BATCH_COUNT)), - lance_datagen::RoundingBehavior::RoundUp, - ) - .map_err(|e| { - pyo3::exceptions::PyValueError::new_err(format!("Failed to generate batches: {}", e)) - })? + if rows_in_batch.is_some() && bytes_in_batch.is_some() { + return Err(pyo3::exceptions::PyValueError::new_err( + "rows_in_batch and bytes_in_batch are mutually exclusive", + )); + } + let builder = lance_datagen::rand(&schema.0); + let batch_count = BatchCount::from(batch_count.unwrap_or(DEFAULT_BATCH_COUNT)); + let reader: Box = match rows_in_batch { + Some(rows) => Box::new(builder.into_reader_rows(RowCount::from(rows), batch_count)), + None => Box::new( + builder + .into_reader_bytes( + ByteCount::from(bytes_in_batch.unwrap_or(DEFAULT_BATCH_SIZE_BYTES)), + batch_count, + lance_datagen::RoundingBehavior::RoundUp, + ) + .map_err(|e| { + pyo3::exceptions::PyValueError::new_err(format!( + "Failed to generate batches: {}", + e + )) + })?, + ), + }; + reader .map(|item| { item.map(PyArrowType::from).map_err(|e| { pyo3::exceptions::PyValueError::new_err(format!("Failed to generate batch: {}", e))