From b198e39ec741902a94d635d3bd7427ffe3afcf37 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Tue, 28 Jul 2026 11:39:35 -0700 Subject: [PATCH] test(python): add merge_insert benchmarks to ci_benchmarks The ci_benchmarks suite had no merge_insert coverage, so there was no way to see how the legacy indexed path compares to the DataFusion path, or to catch a regression in either. Adds 111 parametrized benchmarks covering the source/target ratio sweep, key distribution, cold vs warm index cache, partial-column write amplification, clause shapes, target layouts, and streaming sources. Every benchmark that can run on both paths is parametrized on `use_index`, which is the only knob that selects between them today. Also adds a `setup=` hook to the `io_mem_benchmark` fixture, which previously had no way to reset state between runs -- required for any benchmark that mutates its dataset. Co-Authored-By: Claude Opus 5 (1M context) --- python/python/ci_benchmarks/benchmark.py | 36 +- .../benchmarks/test_merge_insert.py | 766 ++++++++++++++++++ .../python/ci_benchmarks/datagen/gen_all.py | 4 + .../ci_benchmarks/datagen/merge_insert.py | 394 +++++++++ 4 files changed, 1194 insertions(+), 6 deletions(-) create mode 100644 python/python/ci_benchmarks/benchmarks/test_merge_insert.py create mode 100644 python/python/ci_benchmarks/datagen/merge_insert.py diff --git a/python/python/ci_benchmarks/benchmark.py b/python/python/ci_benchmarks/benchmark.py index 7d80596e305..8d17bef5ae5 100644 --- a/python/python/ci_benchmarks/benchmark.py +++ b/python/python/ci_benchmarks/benchmark.py @@ -20,7 +20,7 @@ def workload(dataset): import json from dataclasses import dataclass -from typing import Any, Callable, List +from typing import Any, Callable, List, Optional import pytest @@ -90,6 +90,7 @@ def __call__( func: Callable, dataset: Any, warmup: bool = True, + setup: Optional[Callable[[], Any]] = None, ) -> Any: """ Run a benchmark function with IO and memory tracking. @@ -102,28 +103,47 @@ def __call__( The dataset to pass to the function. warmup : bool, default True Whether to run a warmup iteration before measuring. + setup : Callable, optional + Called before the warmup run and again before the measured run. + Required for workloads that mutate the dataset, so each run starts + from the same state. Setup runs before IO stats are reset, so its + own IO is never attributed to the benchmark. If it returns a value, + that value is passed to ``func`` instead of ``dataset`` -- use this + when the reset produces a new dataset handle. Returns ------- Any The return value of the benchmark function. """ + + def prepare() -> Any: + if setup is None: + return dataset + replacement = setup() + return dataset if replacement is None else replacement + # Warmup run (not measured) if warmup: - func(dataset) + func(prepare()) + + target = prepare() - # Reset IO stats before the measured run + # Reset IO stats before the measured run. Stats are tracked on the + # ObjectStore, which is shared by every handle from the same session, + # so resetting/reading through `dataset` also captures work that `func` + # performs through a handle returned by `setup`. dataset.io_stats_incremental() # Run with memory tracking if available if MEMTEST_AVAILABLE: memtest.reset_stats() - result = func(dataset) + result = func(target) mem_stats = memtest.get_stats() self._stats.peak_bytes = mem_stats["peak_bytes"] self._stats.total_allocations = mem_stats["total_allocations"] else: - result = func(dataset) + result = func(target) # Capture IO stats io_stats = dataset.io_stats_incremental() @@ -159,7 +179,11 @@ def workload(dataset): if marker is None: # Not an io_memory_benchmark test, return a simple passthrough class PassthroughBenchmark: - def __call__(self, func, dataset, warmup=True): + def __call__(self, func, dataset, warmup=True, setup=None): + if setup is not None: + replacement = setup() + if replacement is not None: + dataset = replacement return func(dataset) yield PassthroughBenchmark() diff --git a/python/python/ci_benchmarks/benchmarks/test_merge_insert.py b/python/python/ci_benchmarks/benchmarks/test_merge_insert.py new file mode 100644 index 00000000000..b1d0170db24 --- /dev/null +++ b/python/python/ci_benchmarks/benchmarks/test_merge_insert.py @@ -0,0 +1,766 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The Lance Authors + +"""Benchmarks for `merge_insert`. + +merge_insert has two execution paths and the routing between them is +structural, not a flag. `use_index` is the one knob that selects between them, +so every benchmark here is parametrized on it: + + ``v1_indexed`` — ``use_index(True)`` with an indexed key. Takes the legacy + indexed-scan path (``create_indexed_scan_joined_stream``). + ``v2_hash`` — ``use_index(False)``. Disables the index gate in + ``can_use_create_plan``, so the DataFusion path + (``LanceRead + HashJoin``) runs instead. + +For a partial-schema source the same knob also selects the write sink: v1 +patches columns in place (``UpdateMode::RewriteColumns``) while v2 rewrites +whole rows (``RewriteRows``). That makes ``write_bytes`` the interesting +metric for the ``test_update_*`` benchmarks. + +Targets are mutated, so each measured run is preceded by an untimed restore to +the ``merge_insert_base`` tag written by ``datagen/merge_insert.py``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import timedelta +from typing import TYPE_CHECKING, Callable, Iterable, Optional, Sequence + +import lance +import numpy as np +import pyarrow as pa +import pytest +from ci_benchmarks.datagen.merge_insert import ( + BASE_TAG, + DELETED_ROW_STRIDE, + FRAGS_NUM_ROWS, + FRAGS_SCHEMA, + NARROW_NUM_ROWS, + NARROW_SCHEMA, + UNINDEXED_TAIL_ROWS, + WIDE_NUM_ROWS, + WIDE_ROWS_PER_FRAGMENT, + WIDE_SCALAR_COLUMNS, + WIDE_SCHEMA, + WIDE_VECTOR_DIM, + narrow_batch, +) +from ci_benchmarks.datasets import get_dataset_uri + +if TYPE_CHECKING: + from lance.dataset import ExecuteResult + +PLANS = ["v1_indexed", "v2_hash"] + +# Brackets the cold-random break-even, which the design analysis puts at +# roughly target_rows / 4096 -- about 2.4K rows for the 10M-row narrow target. +SOURCE_SIZES = [1_000, 10_000, 100_000] + +NARROW_KEYS = ["id_int", "id_uuid7", "id_uuid4"] + + +# --------------------------------------------------------------------------- +# Target management +# --------------------------------------------------------------------------- + + +@dataclass +class Target: + """A merge_insert target that can be rewound to a pristine state. + + ``reset`` returns a handle sitting on a fresh version whose contents match + the ``merge_insert_base`` tag. It is called from an untimed benchmark + setup hook, so the restore never lands in the measurement. + """ + + uri: str + dataset: lance.LanceDataset + base_version: int + + def reset(self, cold: bool = False) -> lance.LanceDataset: + if cold: + # A new session means an empty index cache. `checkout_version` + # deliberately shares the cache, so a cold run cannot reuse + # `self.dataset`. + handle = lance.dataset(self.uri) + else: + handle = self.dataset + reverted = handle.checkout_version(self.base_version) + reverted.restore() + return reverted + + +def _open_target(name: str) -> Iterable[Target]: + uri = get_dataset_uri(name) + dataset = lance.dataset(uri) + base_version = dataset.tags.get_version(BASE_TAG) + if base_version is None: + pytest.skip( + f"Dataset {name} has no {BASE_TAG} tag; " + "run python/ci_benchmarks/datagen/gen_all.py" + ) + + yield Target(uri=uri, dataset=dataset, base_version=base_version) + + # Leave the dataset pristine and drop the versions and data files the + # benchmarks produced. Cleanup never deletes a tagged version, so the tag + # has to move onto the restored version first or the old base would be + # retained forever. + final = dataset.checkout_version(base_version) + final.restore() + final.tags.update(BASE_TAG, final.version) + final.cleanup_old_versions(older_than=timedelta(0), delete_unverified=True) + + +@pytest.fixture(scope="module") +def narrow() -> Iterable[Target]: + yield from _open_target("merge_insert_narrow") + + +@pytest.fixture(scope="module") +def wide() -> Iterable[Target]: + yield from _open_target("merge_insert_wide") + + +@pytest.fixture(scope="module") +def frags() -> Iterable[Target]: + yield from _open_target("merge_insert_frags") + + +@pytest.fixture(scope="module") +def deleted() -> Iterable[Target]: + yield from _open_target("merge_insert_deleted") + + +@pytest.fixture(scope="module") +def unindexed_tail() -> Iterable[Target]: + yield from _open_target("merge_insert_unindexed_tail") + + +# --------------------------------------------------------------------------- +# Source construction +# --------------------------------------------------------------------------- +# +# Sources are built from a contiguous run of row indices. Which key column the +# merge joins on then decides whether those keys are clustered or scattered in +# the index's key order: `id_int` and `id_uuid7` are monotonic in the row index, +# `id_uuid4` scrambles it. This keeps one source builder for all key shapes. + + +def narrow_source(row_indices: np.ndarray) -> pa.Table: + """A full-schema narrow source. ``value`` is offset so updates are real.""" + return pa.Table.from_batches([narrow_batch(row_indices, value_offset=1)]) + + +def existing_rows(num_rows: int, offset: int = 0) -> np.ndarray: + return np.arange(offset, offset + num_rows, dtype=np.int64) + + +def new_rows(num_rows: int) -> np.ndarray: + """Row indices past the end of the target, so every key is a new key.""" + return np.arange(NARROW_NUM_ROWS, NARROW_NUM_ROWS + num_rows, dtype=np.int64) + + +def wide_row_indices(fraction: float) -> np.ndarray: + """Row indices covering ``fraction`` of every fragment's rows. + + Spread evenly within each fragment rather than contiguously: that is the + harder case for the in-place column updater, which has to interleave + updated and untouched values. + """ + per_fragment = max(1, round(fraction * WIDE_ROWS_PER_FRAGMENT)) + num_fragments = WIDE_NUM_ROWS // WIDE_ROWS_PER_FRAGMENT + return np.concatenate( + [ + fragment * WIDE_ROWS_PER_FRAGMENT + + np.linspace(0, WIDE_ROWS_PER_FRAGMENT - 1, per_fragment, dtype=np.int64) + for fragment in range(num_fragments) + ] + ) + + +def wide_source(row_indices: np.ndarray, columns: Sequence[str]) -> pa.Table: + """A wide source carrying ``id_int`` plus ``columns``.""" + arrays = [pa.array(row_indices)] + names = ["id_int"] + for column in columns: + names.append(column) + if column == "vec": + values = np.linspace( + 1.0, + 2.0, + num=len(row_indices) * WIDE_VECTOR_DIM, + dtype=np.float32, + ) + arrays.append( + pa.FixedSizeListArray.from_arrays(pa.array(values), WIDE_VECTOR_DIM) + ) + elif WIDE_SCHEMA.field(column).type == pa.string(): + arrays.append( + pa.array([f"updated_{column}_{v}" for v in row_indices], pa.string()) + ) + else: + arrays.append(pa.array(row_indices + 1)) + return pa.table(arrays, schema=pa.schema([WIDE_SCHEMA.field(n) for n in names])) + + +# --------------------------------------------------------------------------- +# Runner +# --------------------------------------------------------------------------- + + +def run( + benchmark, + target: Target, + job: Callable[[lance.LanceDataset], ExecuteResult], + *, + rounds: int = 3, + cold: bool = False, + warmup: bool = True, + expected_rows: Optional[int] = None, +) -> None: + """Benchmark ``job``, restoring the target before every round. + + ``cold`` opens a fresh session for each round so the index cache starts + empty. ``warmup`` runs one unmeasured round first; turn it off for + expensive shapes where a second pass is not worth the wall clock. + + ``expected_rows`` guards against a benchmark that silently stops doing + work: a semantics regression should fail the run, not post a fast time. + Row count is checked instead of the returned stats because it is + independent of which execution path ran. + """ + state: dict = {} + + def setup() -> None: + state["dataset"] = target.reset(cold=cold) + + def bench() -> None: + dataset = state["dataset"] + state["stats"] = job(dataset) + if expected_rows is not None: + state["row_count"] = dataset.count_rows() + + benchmark.pedantic( + bench, + setup=setup, + rounds=rounds, + iterations=1, + # `setup` runs before each warmup round too, so a warmup never leaves + # the target dirty for the measured rounds. + warmup_rounds=1 if warmup and not cold else 0, + ) + + if expected_rows is not None: + assert state["row_count"] == expected_rows, ( + f"expected {expected_rows} rows after merge_insert, " + f"got {state['row_count']} (stats: {state['stats']})" + ) + + +def upsert( + key: str, source: pa.Table, *, use_index: bool +) -> Callable[[lance.LanceDataset], ExecuteResult]: + def job(dataset: lance.LanceDataset) -> ExecuteResult: + return ( + dataset.merge_insert(key) + .when_matched_update_all() + .when_not_matched_insert_all() + .use_index(use_index) + .execute(source) + ) + + return job + + +def uses_index(plan: str) -> bool: + return plan == "v1_indexed" + + +# --------------------------------------------------------------------------- +# A. Cost model core -- merge_insert_narrow, 10M rows +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("plan", PLANS) +def test_upsert_point(benchmark, narrow: Target, plan: str) -> None: + """Single-row upsert latency -- the extreme where a probe should win.""" + source = narrow_source(existing_rows(1)) + run( + benchmark, + narrow, + upsert("id_int", source, use_index=uses_index(plan)), + rounds=5, + expected_rows=NARROW_NUM_ROWS, + ) + + +@pytest.mark.parametrize("plan", PLANS) +@pytest.mark.parametrize("key", NARROW_KEYS) +@pytest.mark.parametrize("num_rows", SOURCE_SIZES) +def test_upsert_ratio_sweep( + benchmark, narrow: Target, num_rows: int, key: str, plan: str +) -> None: + """Source/target ratio sweep against a warm index.""" + source = narrow_source(existing_rows(num_rows)) + run( + benchmark, + narrow, + upsert(key, source, use_index=uses_index(plan)), + expected_rows=NARROW_NUM_ROWS, + ) + + +@pytest.mark.parametrize("plan", PLANS) +@pytest.mark.parametrize("key", ["id_int", "id_uuid4"]) +@pytest.mark.parametrize("num_rows", [1_000, 100_000]) +def test_upsert_cold_ratio_sweep( + benchmark, narrow: Target, num_rows: int, key: str, plan: str +) -> None: + """Same sweep with a cold index cache, where page reads are not free.""" + source = narrow_source(existing_rows(num_rows)) + run( + benchmark, + narrow, + upsert(key, source, use_index=uses_index(plan)), + cold=True, + expected_rows=NARROW_NUM_ROWS, + ) + + +@pytest.mark.parametrize("plan", PLANS) +@pytest.mark.parametrize("num_rows", [10_000, 100_000]) +def test_upsert_all_new_keys( + benchmark, narrow: Target, num_rows: int, plan: str +) -> None: + """Time-ordered ingest: every key is new, so no index page should be read.""" + source = narrow_source(new_rows(num_rows)) + run( + benchmark, + narrow, + upsert("id_uuid7", source, use_index=uses_index(plan)), + expected_rows=NARROW_NUM_ROWS + num_rows, + ) + + +@pytest.mark.parametrize("num_rows", [1_000, 100_000]) +def test_upsert_unindexed_baseline(benchmark, narrow: Target, num_rows: int) -> None: + """Joining on a column with no index at all -- the no-index reference.""" + source = narrow_source(existing_rows(num_rows)) + run( + benchmark, + narrow, + upsert("id_no_index", source, use_index=True), + expected_rows=NARROW_NUM_ROWS, + ) + + +# --------------------------------------------------------------------------- +# B. Write path -- merge_insert_wide, 1M rows +# --------------------------------------------------------------------------- + +# Fractions are of each fragment's rows, which is what decides whether the +# in-place updater interleaves or rewrites a whole column file. +ROW_FRACTIONS = [0.001, 0.01, 0.1, 1.0] +FRACTION_IDS = ["0.1pct", "1pct", "10pct", "100pct"] + +# Isolates the two pressures a partial-column update can exert: field count +# (scalar columns, cheap per field) and byte volume (the vector column). +PROJECTIONS = { + "one_scalar": WIDE_SCALAR_COLUMNS[:1], + "ten_scalars": WIDE_SCALAR_COLUMNS[:10], + "vector": ["vec"], + "vector_and_ten_scalars": ["vec"] + WIDE_SCALAR_COLUMNS[:10], +} + + +def _wide_rounds(fraction: float) -> int: + # A full-fragment update rewrites every column file it touches, which for + # the vector column is ~1 GB per round. + return 1 if fraction == 1.0 else 3 + + +def update_subset( + source: pa.Table, *, use_index: bool +) -> Callable[[lance.LanceDataset], ExecuteResult]: + """Partial-schema update. No insert clause: matched rows only.""" + + def job(dataset: lance.LanceDataset) -> ExecuteResult: + return ( + dataset.merge_insert("id_int") + .when_matched_update_all() + .use_index(use_index) + .execute(source) + ) + + return job + + +@pytest.mark.parametrize("plan", PLANS) +@pytest.mark.parametrize("fraction", ROW_FRACTIONS, ids=FRACTION_IDS) +def test_update_subset_row_fraction( + benchmark, wide: Target, fraction: float, plan: str +) -> None: + """Row-fraction sweep at a fixed, minimal projection. + + Locates the crossover between patching a column in place and rewriting the + whole column file, without byte volume confounding it. + """ + source = wide_source(wide_row_indices(fraction), PROJECTIONS["one_scalar"]) + run( + benchmark, + wide, + update_subset(source, use_index=uses_index(plan)), + rounds=_wide_rounds(fraction), + warmup=fraction != 1.0, + expected_rows=WIDE_NUM_ROWS, + ) + + +@pytest.mark.parametrize("plan", PLANS) +@pytest.mark.parametrize("fraction", [0.01, 1.0], ids=["1pct", "100pct"]) +@pytest.mark.parametrize("projection", list(PROJECTIONS), ids=list(PROJECTIONS)) +def test_update_subset_projection( + benchmark, wide: Target, projection: str, fraction: float, plan: str +) -> None: + """Projection sweep: field count vs byte volume, plus the cross term.""" + source = wide_source(wide_row_indices(fraction), PROJECTIONS[projection]) + run( + benchmark, + wide, + update_subset(source, use_index=uses_index(plan)), + rounds=_wide_rounds(fraction), + warmup=fraction != 1.0, + expected_rows=WIDE_NUM_ROWS, + ) + + +@pytest.mark.parametrize("plan", PLANS) +@pytest.mark.parametrize("fraction", [0.01, 1.0], ids=["1pct", "100pct"]) +def test_upsert_wide_full_schema( + benchmark, wide: Target, fraction: float, plan: str +) -> None: + """Full-schema baseline for the partial-column benchmarks above.""" + row_indices = wide_row_indices(fraction) + source = wide_source( + row_indices, [f.name for f in WIDE_SCHEMA if f.name != "id_int"] + ) + run( + benchmark, + wide, + upsert("id_int", source, use_index=uses_index(plan)), + rounds=_wide_rounds(fraction), + warmup=fraction != 1.0, + expected_rows=WIDE_NUM_ROWS, + ) + + +# --------------------------------------------------------------------------- +# C. Clause shapes -- merge_insert_narrow, 10K-row source +# --------------------------------------------------------------------------- + +CLAUSE_SOURCE_ROWS = 10_000 + + +@pytest.mark.parametrize("plan", PLANS) +def test_insert_if_not_exists(benchmark, narrow: Target, plan: str) -> None: + """Dedup ingest: half the keys already exist, matched rows are untouched. + + The probe only needs to know whether a key exists, so no target payload has + to be read. + """ + half = CLAUSE_SOURCE_ROWS // 2 + row_indices = np.concatenate([existing_rows(half), new_rows(half)]) + source = narrow_source(row_indices) + + def job(dataset: lance.LanceDataset) -> ExecuteResult: + return ( + dataset.merge_insert("id_int") + .when_not_matched_insert_all() + .use_index(uses_index(plan)) + .execute(source) + ) + + run(benchmark, narrow, job, expected_rows=NARROW_NUM_ROWS + half) + + +@pytest.mark.parametrize("plan", PLANS) +def test_update_only(benchmark, narrow: Target, plan: str) -> None: + """No insert clause: unmatched source rows are dropped.""" + half = CLAUSE_SOURCE_ROWS // 2 + row_indices = np.concatenate([existing_rows(half), new_rows(half)]) + source = narrow_source(row_indices) + + def job(dataset: lance.LanceDataset) -> ExecuteResult: + return ( + dataset.merge_insert("id_int") + .when_matched_update_all() + .use_index(uses_index(plan)) + .execute(source) + ) + + run(benchmark, narrow, job, expected_rows=NARROW_NUM_ROWS) + + +def test_delete_by_source(benchmark, narrow: Target) -> None: + """Deleting rows absent from the source requires a full target scan. + + The index gate in `can_use_create_plan` rejects this shape outright, so + there is no v1 variant to compare against. + """ + source = narrow_source(existing_rows(CLAUSE_SOURCE_ROWS)) + + def job(dataset: lance.LanceDataset) -> ExecuteResult: + return ( + dataset.merge_insert("id_int") + .when_matched_update_all() + .when_not_matched_insert_all() + .when_not_matched_by_source_delete() + .execute(source) + ) + + run(benchmark, narrow, job, expected_rows=CLAUSE_SOURCE_ROWS) + + +@pytest.mark.parametrize("plan", PLANS) +def test_conditional_update(benchmark, narrow: Target, plan: str) -> None: + """A condition on target columns forces the target payload to be read.""" + source = narrow_source(existing_rows(CLAUSE_SOURCE_ROWS)) + + def job(dataset: lance.LanceDataset) -> ExecuteResult: + return ( + dataset.merge_insert("id_int") + .when_matched_update_all(condition="source.value > target.value") + .when_not_matched_insert_all() + .use_index(uses_index(plan)) + .execute(source) + ) + + run(benchmark, narrow, job, expected_rows=NARROW_NUM_ROWS) + + +@pytest.mark.parametrize("plan", PLANS) +def test_composite_key_fully_indexed(benchmark, narrow: Target, plan: str) -> None: + """Both key columns indexed: v1 probes each index and AND-folds.""" + source = narrow_source(existing_rows(CLAUSE_SOURCE_ROWS)) + + def job(dataset: lance.LanceDataset) -> ExecuteResult: + return ( + dataset.merge_insert(["composite_a", "composite_b"]) + .when_matched_update_all() + .when_not_matched_insert_all() + .use_index(uses_index(plan)) + .execute(source) + ) + + run(benchmark, narrow, job, expected_rows=NARROW_NUM_ROWS) + + +def test_composite_key_partially_indexed(benchmark, narrow: Target) -> None: + """One key column unindexed, which forces the hash join regardless of flag. + + A partial index probe would under-match, so `can_use_create_plan` keeps + this shape off the indexed path. + """ + source = narrow_source(existing_rows(CLAUSE_SOURCE_ROWS)) + + def job(dataset: lance.LanceDataset) -> ExecuteResult: + # `composite_a` is indexed, `id_no_index` is not. Both are equal to the + # row index in source and target, so every source row matches. + return ( + dataset.merge_insert(["composite_a", "id_no_index"]) + .when_matched_update_all() + .when_not_matched_insert_all() + .execute(source) + ) + + run(benchmark, narrow, job, expected_rows=NARROW_NUM_ROWS) + + +# --------------------------------------------------------------------------- +# D. Target shape +# --------------------------------------------------------------------------- + +TARGET_SHAPE_ROWS = 10_000 + + +@pytest.mark.parametrize("plan", PLANS) +def test_upsert_unindexed_tail(benchmark, unindexed_tail: Target, plan: str) -> None: + """10% of the target was appended after indexing, so a scan is unioned in.""" + source = narrow_source(existing_rows(TARGET_SHAPE_ROWS)) + total_rows = NARROW_NUM_ROWS + UNINDEXED_TAIL_ROWS + run( + benchmark, + unindexed_tail, + upsert("id_int", source, use_index=uses_index(plan)), + expected_rows=total_rows, + ) + + +@pytest.mark.parametrize("plan", PLANS) +def test_upsert_with_deletion_files(benchmark, deleted: Target, plan: str) -> None: + """Every fragment carries a deletion file, which the probe has to mask.""" + source = narrow_source(existing_rows(TARGET_SHAPE_ROWS)) + # Row 0 and every DELETED_ROW_STRIDE-th row after it were deleted at + # generation time, so those source rows insert rather than update. + reinserted = TARGET_SHAPE_ROWS // DELETED_ROW_STRIDE + total_rows = NARROW_NUM_ROWS - NARROW_NUM_ROWS // DELETED_ROW_STRIDE + reinserted + run( + benchmark, + deleted, + upsert("id_int", source, use_index=uses_index(plan)), + expected_rows=total_rows, + ) + + +@pytest.mark.parametrize("plan", PLANS) +def test_upsert_many_small_fragments(benchmark, frags: Target, plan: str) -> None: + """10K fragments of 1K rows, to expose per-fragment overhead.""" + row_indices = existing_rows(TARGET_SHAPE_ROWS) + source = pa.table( + [pa.array(row_indices), pa.array(row_indices + 1)], schema=FRAGS_SCHEMA + ) + run( + benchmark, + frags, + upsert("id_int", source, use_index=uses_index(plan)), + expected_rows=FRAGS_NUM_ROWS, + ) + + +# --------------------------------------------------------------------------- +# E. Memory and regime edges +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("plan", PLANS) +def test_upsert_streaming_source(benchmark, narrow: Target, plan: str) -> None: + """A one-shot reader source, which v1 has to buffer in full to fork it. + + Peak memory is the number of interest here. + """ + num_rows = 1_000_000 + batch_size = 100_000 + + def make_reader() -> pa.RecordBatchReader: + def batches(): + for start in range(0, num_rows, batch_size): + yield narrow_batch( + existing_rows(batch_size, offset=start), value_offset=1 + ) + + return pa.RecordBatchReader.from_batches(NARROW_SCHEMA, batches()) + + def job(dataset: lance.LanceDataset) -> ExecuteResult: + return ( + dataset.merge_insert("id_int") + .when_matched_update_all() + .when_not_matched_insert_all() + .use_index(uses_index(plan)) + .execute(make_reader()) + ) + + run( + benchmark, + narrow, + job, + rounds=1, + warmup=False, + expected_rows=NARROW_NUM_ROWS, + ) + + +def test_upsert_source_equals_target(benchmark, narrow: Target) -> None: + """Source the same size as the target -- the probe must not be chosen here. + + v2 only: the v1 indexed path cannot run this shape at all. Its source-side + hash join asks for more than the whole memory pool and the operation fails + with "Resources exhausted". So there is no v1 baseline to compare against, + and this benchmark exists to keep the v2 path honest at this size. + """ + source = narrow_source(existing_rows(NARROW_NUM_ROWS)) + run( + benchmark, + narrow, + upsert("id_int", source, use_index=False), + rounds=1, + warmup=False, + expected_rows=NARROW_NUM_ROWS, + ) + + +# --------------------------------------------------------------------------- +# IO / memory variants +# --------------------------------------------------------------------------- +# +# A subset of the above, re-run under the io_memory_benchmark fixture. Write +# amplification (`write_bytes`) is the headline for the partial-column cases and +# peak memory is the headline for the streaming source. + + +@pytest.mark.io_memory_benchmark() +@pytest.mark.parametrize("plan", PLANS) +@pytest.mark.parametrize("key", NARROW_KEYS) +@pytest.mark.parametrize("num_rows", [1_000, 100_000]) +def test_io_mem_upsert_ratio( + io_mem_benchmark, narrow: Target, num_rows: int, key: str, plan: str +) -> None: + source = narrow_source(existing_rows(num_rows)) + job = upsert(key, source, use_index=uses_index(plan)) + io_mem_benchmark(job, narrow.dataset, setup=lambda: narrow.reset()) + + +@pytest.mark.io_memory_benchmark() +@pytest.mark.parametrize("plan", PLANS) +@pytest.mark.parametrize("fraction", ROW_FRACTIONS, ids=FRACTION_IDS) +def test_io_mem_update_subset_row_fraction( + io_mem_benchmark, wide: Target, fraction: float, plan: str +) -> None: + source = wide_source(wide_row_indices(fraction), PROJECTIONS["one_scalar"]) + job = update_subset(source, use_index=uses_index(plan)) + io_mem_benchmark( + job, + wide.dataset, + warmup=fraction != 1.0, + setup=lambda: wide.reset(), + ) + + +@pytest.mark.io_memory_benchmark() +@pytest.mark.parametrize("plan", PLANS) +@pytest.mark.parametrize("projection", list(PROJECTIONS), ids=list(PROJECTIONS)) +def test_io_mem_update_subset_projection( + io_mem_benchmark, wide: Target, projection: str, plan: str +) -> None: + source = wide_source(wide_row_indices(0.01), PROJECTIONS[projection]) + job = update_subset(source, use_index=uses_index(plan)) + io_mem_benchmark(job, wide.dataset, setup=lambda: wide.reset()) + + +@pytest.mark.io_memory_benchmark() +@pytest.mark.parametrize("plan", PLANS) +def test_io_mem_upsert_streaming_source( + io_mem_benchmark, narrow: Target, plan: str +) -> None: + num_rows = 1_000_000 + batch_size = 100_000 + + def job(dataset: lance.LanceDataset) -> ExecuteResult: + def batches(): + for start in range(0, num_rows, batch_size): + yield narrow_batch( + existing_rows(batch_size, offset=start), value_offset=1 + ) + + reader = pa.RecordBatchReader.from_batches(NARROW_SCHEMA, batches()) + return ( + dataset.merge_insert("id_int") + .when_matched_update_all() + .when_not_matched_insert_all() + .use_index(uses_index(plan)) + .execute(reader) + ) + + io_mem_benchmark(job, narrow.dataset, warmup=False, setup=lambda: narrow.reset()) diff --git a/python/python/ci_benchmarks/datagen/gen_all.py b/python/python/ci_benchmarks/datagen/gen_all.py index d5120d20ff7..79725708bb8 100644 --- a/python/python/ci_benchmarks/datagen/gen_all.py +++ b/python/python/ci_benchmarks/datagen/gen_all.py @@ -8,6 +8,7 @@ from ci_benchmarks.datagen.basic import gen_basic from ci_benchmarks.datagen.count_rows import gen_count_rows from ci_benchmarks.datagen.lineitems import gen_tcph +from ci_benchmarks.datagen.merge_insert import gen_merge_insert from ci_benchmarks.datagen.wikipedia import gen_wikipedia @@ -44,6 +45,9 @@ def setup_logging(): LOGGER.info("Generating count_rows benchmark dataset...") gen_count_rows() + LOGGER.info("Generating merge_insert benchmark datasets...") + gen_merge_insert() + LOGGER.info("=" * 80) LOGGER.info("All datasets generated successfully!") LOGGER.info("=" * 80) diff --git a/python/python/ci_benchmarks/datagen/merge_insert.py b/python/python/ci_benchmarks/datagen/merge_insert.py new file mode 100644 index 00000000000..ebda5730aa2 --- /dev/null +++ b/python/python/ci_benchmarks/datagen/merge_insert.py @@ -0,0 +1,394 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The Lance Authors + +"""Generate the merge_insert benchmark datasets. + +merge_insert benchmarks mutate their target, so each dataset carries a +``merge_insert_base`` tag pointing at a pristine version. Benchmarks restore +to that tag before every measured run (see ``benchmarks/test_merge_insert.py``), +which makes the suite tolerant of a crashed run leaving stale versions behind. + +Datasets +-------- +``merge_insert_narrow`` + 10M rows, 10 fragments. The key columns differ along two independent + axes so benchmarks can separate them: + + * ``id_int`` vs ``id_uuid7`` -- key *type* (int64 vs string), both + clustered. + * ``id_uuid7`` vs ``id_uuid4`` -- key *distribution* (clustered vs + random), both string. + + ``id_no_index`` holds the same values as ``id_int`` with no index, for the + "user never built an index" baseline. ``composite_a``/``composite_b`` are + both indexed, for composite-key probes. + +``merge_insert_wide`` + 1M rows, 10 fragments. 20 narrow scalar columns (~200 MB) plus one + 256-dim float32 vector column (~1 GB). These exert different pressures on + a partial-column update: the scalar columns stress field count and the + number of distinct buffers to schedule, the vector column stresses raw + byte volume. + +``merge_insert_frags`` + 10M rows in 10K fragments of 1K rows, for per-fragment overhead. + +``merge_insert_deleted`` + ``merge_insert_narrow`` layout with a deletion file on every fragment. + +``merge_insert_unindexed_tail`` + ``merge_insert_narrow`` layout where 10% of the rows were appended after + the index was built, so the probe must union an unindexed scan. +""" + +from __future__ import annotations + +import lance +import numpy as np +import pyarrow as pa +from lance.log import LOGGER + +from ci_benchmarks.datasets import get_dataset_uri + +# Tag marking the pristine version that benchmarks restore to. +BASE_TAG = "merge_insert_base" + +NARROW_NUM_ROWS = 10_000_000 +NARROW_ROWS_PER_FRAGMENT = 1_000_000 + +WIDE_NUM_ROWS = 1_000_000 +WIDE_ROWS_PER_FRAGMENT = 100_000 +WIDE_NUM_SCALAR_COLUMNS = 20 +WIDE_VECTOR_DIM = 256 + +FRAGS_NUM_ROWS = 10_000_000 +FRAGS_ROWS_PER_FRAGMENT = 1_000 + +# Rows appended to `merge_insert_unindexed_tail` after the index is built. +UNINDEXED_TAIL_ROWS = NARROW_NUM_ROWS // 10 + +# `merge_insert_deleted` deletes every Nth row, which puts a deletion file on +# every fragment. +DELETED_ROW_STRIDE = 1000 + +_BATCH_SIZE = 1_000_000 + +NARROW_INDEXED_COLUMNS = [ + "id_int", + "id_uuid7", + "id_uuid4", + "composite_a", + "composite_b", +] + +NARROW_SCHEMA = pa.schema( + [ + ("id_int", pa.int64()), + ("id_uuid7", pa.string()), + ("id_uuid4", pa.string()), + ("id_no_index", pa.int64()), + ("composite_a", pa.int64()), + ("composite_b", pa.int64()), + ("value", pa.int64()), + ] +) + +FRAGS_SCHEMA = pa.schema([("id_int", pa.int64()), ("value", pa.int64())]) + +WIDE_SCALAR_COLUMNS = [f"scalar_{i}" for i in range(WIDE_NUM_SCALAR_COLUMNS)] + +# Half int64, half string, so a partial-column update touches a mix of +# fixed-width and variable-width buffers. +WIDE_SCHEMA = pa.schema( + [("id_int", pa.int64())] + + [ + (name, pa.int64() if i % 2 == 0 else pa.string()) + for i, name in enumerate(WIDE_SCALAR_COLUMNS) + ] + + [("vec", pa.list_(pa.float32(), WIDE_VECTOR_DIM))] +) + + +_GOLDEN = np.uint64(0x9E3779B97F4A7C15) +_MIX1 = np.uint64(0xBF58476D1CE4E5B9) +_MIX2 = np.uint64(0x94D049BB133111EB) +_HEX_DIGITS = np.frombuffer(b"0123456789abcdef", dtype="S1") +_NIBBLE_SHIFTS = np.arange(60, -4, -4, dtype=np.uint64) +# Two uint64 halves rendered as hex. +_KEY_WIDTH = 32 + + +def _scramble(values: np.ndarray) -> np.ndarray: + """splitmix64 finalizer. Wrapping uint64 arithmetic, vectorized.""" + x = values.astype(np.uint64) * _GOLDEN + x = (x ^ (x >> np.uint64(30))) * _MIX1 + x = (x ^ (x >> np.uint64(27))) * _MIX2 + return x ^ (x >> np.uint64(31)) + + +def _hex_keys(high: np.ndarray, low: np.ndarray) -> pa.Array: + """Format two uint64 columns as 32-character lowercase hex strings. + + Vectorized because the narrow dataset needs 10M of these and the + benchmarks rebuild source keys from row indices on every run. + """ + # One nibble at a time into a preallocated byte matrix. Broadcasting all 16 + # shifts at once would materialize an (n, 16) uint64 array per half, which is + # 2.5 GB of scratch for a 10M-row source. + num_rows = len(high) + half = len(_NIBBLE_SHIFTS) + digits = np.empty((num_rows, _KEY_WIDTH), dtype="S1") + for position, shift in enumerate(_NIBBLE_SHIFTS): + digits[:, position] = _HEX_DIGITS[(high >> shift) & np.uint64(0xF)] + digits[:, position + half] = _HEX_DIGITS[(low >> shift) & np.uint64(0xF)] + + # Built from buffers rather than `pa.array`, which splits a numpy byte array + # into chunks above ~1M elements. RecordBatch needs a contiguous Array. + offsets = np.arange(0, _KEY_WIDTH * (num_rows + 1), _KEY_WIDTH, dtype=np.int32) + return pa.StringArray.from_buffers( + num_rows, pa.py_buffer(offsets), pa.py_buffer(digits) + ) + + +def uuid7_keys(row_indices: np.ndarray) -> pa.Array: + """Sortable, UUIDv7-shaped keys: monotonic prefix, scrambled suffix. + + Real UUIDv7 keys are time-ordered, so a stream of them lands in a narrow, + advancing slice of the index. Reproducing that ordering is what matters + for the benchmark; the exact bit layout is not. + + Deterministic in the row index so benchmarks can reconstruct a key without + reading the dataset. + """ + row_indices = row_indices.astype(np.uint64) + return _hex_keys(row_indices, _scramble(row_indices)) + + +def uuid4_keys(row_indices: np.ndarray) -> pa.Array: + """Same width as `uuid7_keys`, but scattered across the index key space. + + Deterministic for the same reason, but the leading bytes are scrambled, so + a contiguous run of row indices maps to keys spread over the whole index + rather than a narrow slice. + """ + row_indices = row_indices.astype(np.uint64) + return _hex_keys(_scramble(row_indices), row_indices) + + +def narrow_batch(row_indices: np.ndarray, value_offset: int = 0) -> pa.RecordBatch: + """Build a full-schema narrow batch from row indices. + + Shared with the benchmarks, which reconstruct source rows from row indices. + ``value_offset`` shifts the payload column so an update writes a value that + differs from what the target already holds. + """ + return pa.record_batch( + [ + pa.array(row_indices), + uuid7_keys(row_indices), + uuid4_keys(row_indices), + pa.array(row_indices), + pa.array(row_indices), + # `composite_b` is a deterministic function of `composite_a` so a + # source row can target an existing composite key without the + # benchmark tracking extra state. + pa.array(row_indices % 1024), + pa.array(row_indices + value_offset), + ], + schema=NARROW_SCHEMA, + ) + + +def _narrow_data(num_rows: int, offset: int = 0): + LOGGER.info("Generating %d narrow rows starting at %d", num_rows, offset) + for start in range(offset, offset + num_rows, _BATCH_SIZE): + count = min(_BATCH_SIZE, offset + num_rows - start) + yield narrow_batch(np.arange(start, start + count, dtype=np.int64)) + + +def _frags_data(num_rows: int): + LOGGER.info("Generating %d small-fragment rows", num_rows) + for start in range(0, num_rows, _BATCH_SIZE): + ids = np.arange( + start, start + min(_BATCH_SIZE, num_rows - start), dtype=np.int64 + ) + yield pa.record_batch([pa.array(ids), pa.array(ids)], schema=FRAGS_SCHEMA) + + +def _wide_batch(offset: int, num_rows: int) -> pa.RecordBatch: + ids = np.arange(offset, offset + num_rows, dtype=np.int64) + columns: list[pa.Array] = [pa.array(ids)] + for i in range(WIDE_NUM_SCALAR_COLUMNS): + if i % 2 == 0: + columns.append(pa.array(ids + i)) + else: + columns.append(pa.array([f"s{i}_{v}" for v in ids], type=pa.string())) + # Deterministic float payload; the values are irrelevant, the byte volume + # is the point. + vectors = np.linspace( + 0.0, 1.0, num=num_rows * WIDE_VECTOR_DIM, dtype=np.float32 + ).reshape(num_rows, WIDE_VECTOR_DIM) + columns.append( + pa.FixedSizeListArray.from_arrays( + pa.array(vectors.reshape(-1)), WIDE_VECTOR_DIM + ) + ) + return pa.record_batch(columns, schema=WIDE_SCHEMA) + + +def _wide_data(num_rows: int): + LOGGER.info("Generating %d wide rows", num_rows) + # The vector column makes full batches large, so use a smaller batch here. + batch_size = WIDE_ROWS_PER_FRAGMENT + for start in range(0, num_rows, batch_size): + yield _wide_batch(start, min(batch_size, num_rows - start)) + + +def _tag_base(ds: lance.LanceDataset) -> None: + """Point BASE_TAG at the current version, creating the tag if needed.""" + if BASE_TAG in ds.tags.list(): + ds.tags.update(BASE_TAG, ds.version) + else: + ds.tags.create(BASE_TAG, ds.version) + + +def _already_generated(uri: str, expected_rows: int) -> bool: + """True when a usable dataset with a pristine base tag already exists. + + A previous benchmark run may have left extra versions behind, so the row + count is checked at the tagged version rather than at the latest one. + """ + try: + ds = lance.dataset(uri) + except ValueError: + return False + base_version = ds.tags.get_version(BASE_TAG) + if base_version is None: + return False + return ds.checkout_version(base_version).count_rows() == expected_rows + + +def _gen( + name: str, + data, + schema: pa.Schema, + expected_rows: int, + rows_per_fragment: int, + indexed_columns: list[str], +) -> lance.LanceDataset: + dataset_uri = get_dataset_uri(name) + if _already_generated(dataset_uri, expected_rows): + LOGGER.info("Dataset %s already exists, skipping", name) + return lance.dataset(dataset_uri) + + LOGGER.info("Creating dataset %s", name) + ds = lance.write_dataset( + data, + dataset_uri, + schema=schema, + mode="overwrite", + max_rows_per_file=rows_per_fragment, + max_rows_per_group=min(rows_per_fragment, 100_000), + ) + for column in indexed_columns: + LOGGER.info("Building BTREE index on %s.%s", name, column) + ds.create_scalar_index(column, "BTREE") + _tag_base(ds) + return ds + + +def gen_merge_insert_narrow() -> lance.LanceDataset: + return _gen( + "merge_insert_narrow", + _narrow_data(NARROW_NUM_ROWS), + NARROW_SCHEMA, + NARROW_NUM_ROWS, + NARROW_ROWS_PER_FRAGMENT, + NARROW_INDEXED_COLUMNS, + ) + + +def gen_merge_insert_wide() -> lance.LanceDataset: + return _gen( + "merge_insert_wide", + _wide_data(WIDE_NUM_ROWS), + WIDE_SCHEMA, + WIDE_NUM_ROWS, + WIDE_ROWS_PER_FRAGMENT, + ["id_int"], + ) + + +def gen_merge_insert_frags() -> lance.LanceDataset: + return _gen( + "merge_insert_frags", + _frags_data(FRAGS_NUM_ROWS), + FRAGS_SCHEMA, + FRAGS_NUM_ROWS, + FRAGS_ROWS_PER_FRAGMENT, + ["id_int"], + ) + + +def gen_merge_insert_deleted() -> lance.LanceDataset: + """Narrow layout with a deletion file on every fragment. + + One row in every `DELETED_ROW_STRIDE` is deleted, which lands in every + fragment. + """ + name = "merge_insert_deleted" + dataset_uri = get_dataset_uri(name) + expected_rows = NARROW_NUM_ROWS - NARROW_NUM_ROWS // DELETED_ROW_STRIDE + if _already_generated(dataset_uri, expected_rows): + LOGGER.info("Dataset %s already exists, skipping", name) + return lance.dataset(dataset_uri) + + ds = _gen( + name, + _narrow_data(NARROW_NUM_ROWS), + NARROW_SCHEMA, + NARROW_NUM_ROWS, + NARROW_ROWS_PER_FRAGMENT, + NARROW_INDEXED_COLUMNS, + ) + LOGGER.info("Deleting rows from %s to create deletion files", name) + ds.delete(f"id_int % {DELETED_ROW_STRIDE} == 0") + _tag_base(ds) + return ds + + +def gen_merge_insert_unindexed_tail() -> lance.LanceDataset: + """Narrow layout where the last 10% of rows are not covered by the index.""" + name = "merge_insert_unindexed_tail" + dataset_uri = get_dataset_uri(name) + expected_rows = NARROW_NUM_ROWS + UNINDEXED_TAIL_ROWS + if _already_generated(dataset_uri, expected_rows): + LOGGER.info("Dataset %s already exists, skipping", name) + return lance.dataset(dataset_uri) + + _gen( + name, + _narrow_data(NARROW_NUM_ROWS), + NARROW_SCHEMA, + NARROW_NUM_ROWS, + NARROW_ROWS_PER_FRAGMENT, + NARROW_INDEXED_COLUMNS, + ) + LOGGER.info("Appending %d unindexed rows to %s", UNINDEXED_TAIL_ROWS, name) + ds = lance.write_dataset( + _narrow_data(UNINDEXED_TAIL_ROWS, offset=NARROW_NUM_ROWS), + dataset_uri, + schema=NARROW_SCHEMA, + mode="append", + max_rows_per_file=NARROW_ROWS_PER_FRAGMENT, + ) + _tag_base(ds) + return ds + + +def gen_merge_insert() -> None: + gen_merge_insert_narrow() + gen_merge_insert_wide() + gen_merge_insert_frags() + gen_merge_insert_deleted() + gen_merge_insert_unindexed_tail()