test(python): add data overlay benchmark suite

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)

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)

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))

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)