From b198e39ec741902a94d635d3bd7427ffe3afcf37 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Tue, 28 Jul 2026 11:39:35 -0700 Subject: [PATCH 1/3] 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() From 90b08554dfa550f7cf6ac4f85e5b0b7614f2cf38 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Tue, 28 Jul 2026 14:13:37 -0700 Subject: [PATCH 2/3] feat(merge_insert): probe scalar indices on the DataFusion plan path The DataFusion merge path reached matched target rows by scanning the whole table and hash-joining against the source. Only the legacy path could probe a scalar index on the merge keys, so any merge that needed the DataFusion path scanned the full target even when an index would have found the same rows. Adds an index-probe target provider: matched rows come from an `IsIn` probe of every indexed merge key, taken by row address, unioned with a scan of the fragments no chosen index covers. It reports the same schema as the full-scan provider, so the join, action expression, and write sink are unchanged. Routing is unchanged: this only affects merges that already used the DataFusion path and have an indexed key, chiefly a composite key with some columns indexed. `explain_plan` now uses a re-scannable source so the probe is visible, and `analyze_plan` builds its provider the way execution does. Co-Authored-By: Claude Opus 5 (1M context) --- python/python/lance/dataset.py | 11 +- rust/lance/src/dataset/write/merge_insert.rs | 433 ++++++++++++++++-- .../src/dataset/write/merge_insert/probe.rs | 217 +++++++++ 3 files changed, 607 insertions(+), 54 deletions(-) create mode 100644 rust/lance/src/dataset/write/merge_insert/probe.rs diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 8587ed3badf..37aa0018950 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -640,11 +640,11 @@ def explain_plan( MergeInsert: on=[id], when_matched=UpdateAll, when_not_matched=InsertAll, ... CoalescePartitionsExec ProjectionExec: expr=[...] - HashJoinExec: mode=CollectLeft, join_type=Right, ... - LanceRead: uri=test_dataset/data, projection=[id], ... - RepartitionExec: ... + RepartitionExec: ... + HashJoinExec: mode=CollectLeft, join_type=Left, ... ProjectionExec: expr=[..., true as __merge_source_sentinel] - StreamingTableExec: partition_sizes=1, ... + DataSourceExec: partitions=1, partition_sizes=[1] + LanceRead: uri=test_dataset/data, projection=[id], ... >>> # Or with explicit schema @@ -658,7 +658,8 @@ def explain_plan( MergeInsert: on=[id], when_matched=UpdateAll, when_not_matched=InsertAll, ... CoalescePartitionsExec ProjectionExec: expr=[...] - HashJoinExec: mode=CollectLeft, join_type=Right, ... + RepartitionExec: ... + HashJoinExec: mode=CollectLeft, join_type=Left, ... ... """ return super(MergeInsertBuilder, self).explain_plan(schema, verbose=verbose) diff --git a/rust/lance/src/dataset/write/merge_insert.rs b/rust/lance/src/dataset/write/merge_insert.rs index a325d5de037..57d57eb0d3a 100644 --- a/rust/lance/src/dataset/write/merge_insert.rs +++ b/rust/lance/src/dataset/write/merge_insert.rs @@ -147,6 +147,7 @@ use tracing::error; mod assign_action; mod exec; mod logical_plan; +mod probe; // "update if" expressions typically compare fields from the source table to the target table. // These tables have the same schema and so filter expressions need to differentiate. To do that @@ -1711,7 +1712,7 @@ impl MergeInsertJob { batches: Vec, ) -> Result { let provider = self.batches_to_provider(batches)?; - self.execute_uncommitted_impl(provider).await + self.execute_uncommitted_impl(provider, true).await } /// Wrap materialized batches in a multi-partition in-memory [`MemTable`]. @@ -1783,6 +1784,7 @@ impl MergeInsertJob { let wrapper = MergeInsertJobWithProvider { job: self, provider, + replayable, attempt_count: Arc::new(AtomicU32::new(0)), }; @@ -1797,7 +1799,7 @@ impl MergeInsertJob { source: impl StreamingWriteSource, ) -> Result { let stream = source.into_stream(); - self.execute_uncommitted_impl(one_shot_provider(stream)?) + self.execute_uncommitted_impl(one_shot_provider(stream)?, false) .await } @@ -1816,13 +1818,53 @@ impl MergeInsertJob { } } - async fn create_plan(self, provider: Arc) -> Result> { - // Goal: we shouldn't manually have to specify which columns to scan. - // DataFusion's optimizer should be able to automatically perform - // projection pushdown for us. - // Goal: we shouldn't have to add new branches in this code to handle - // indexed vs non-indexed cases. That should be handled by optimizer rules. - let session_ctx = SessionContext::new(); + /// The scalar-index lookups the target scan should probe instead of scanning + /// the whole table, or `None` to scan. + /// + /// The probe reaches only rows whose keys appear in the source, so it can + /// only replace the scan when the merge cannot touch a row the source never + /// mentions. It also scans the source a second time to collect key values, + /// so the source must be re-scannable. + /// + /// Probing an indexed subset of a composite key is allowed: the probe + /// over-matches and the join applies the real key predicate. Under-matching + /// would not be allowed, which is why the caller unions in the fragments no + /// chosen index covers. + async fn probe_lookups( + &self, + source_replayable: bool, + ) -> Result>> { + let eligible = source_replayable + // Honor an explicit opt-out from consulting an index. + && self.params.use_index + // Reading rows from a probe's output requires the v2 storage format. + && !self.dataset.is_legacy_storage() + // Deleting unmatched target rows needs rows the probe never reaches. + && matches!( + self.params.delete_not_matched_by_source, + WhenNotMatchedBySource::Keep + ); + if !eligible { + return Ok(None); + } + + let indexed = self.indexed_join_keys().await?; + if indexed.is_empty() { + return Ok(None); + } + Ok(Some(indexed)) + } + + /// Build the target-side [`TableProvider`], probing a scalar index where that + /// is possible and scanning the whole table otherwise. + /// + /// Both providers report the same schema, so the rest of the plan does not + /// change with the choice. + async fn target_provider( + &self, + source: &Arc, + source_replayable: bool, + ) -> Result> { let binary_blob_field_ids = self .dataset .schema() @@ -1830,17 +1872,53 @@ impl MergeInsertJob { .filter(|field| field.is_blob() && !field.is_blob_v2()) .map(|field| field.id as u32) .collect(); - let target_provider = Arc::new( - crate::datafusion::dataframe::LanceTableProvider::new_with_ordering( - self.dataset.clone(), - true, - true, - false, - ) - .with_blob_handling(lance_core::datatypes::BlobHandling::SomeBlobsBinary( - binary_blob_field_ids, + let blob_handling = + lance_core::datatypes::BlobHandling::SomeBlobsBinary(binary_blob_field_ids); + + match self.probe_lookups(source_replayable).await? { + Some(indexed) => { + let unindexed_fragments = self.unindexed_fragments_for_keys(&indexed).await?; + let lookups = indexed + .into_iter() + .map(|(column, index)| IndexLookup::new(column, index.name)) + .collect(); + Ok(Arc::new(probe::IndexProbeTarget::try_new( + self.dataset.clone(), + source.clone(), + lookups, + unindexed_fragments, + blob_handling, + )?)) + } + None => Ok(Arc::new( + crate::datafusion::dataframe::LanceTableProvider::new_with_ordering( + self.dataset.clone(), + true, + true, + false, + ) + .with_blob_handling(blob_handling), )), - ); + } + } + + /// Build the physical plan for the merge. + /// + /// `source_replayable` states whether `provider` can be scanned more than + /// once; the index-probe target scan needs a second scan of the source and is + /// not used when it cannot have one. + async fn create_plan( + self, + provider: Arc, + source_replayable: bool, + ) -> Result> { + // Goal: we shouldn't manually have to specify which columns to scan. + // DataFusion's optimizer should be able to automatically perform + // projection pushdown for us. + // Goal: we shouldn't have to add new branches in this code to handle + // indexed vs non-indexed cases. That should be handled by optimizer rules. + let session_ctx = SessionContext::new(); + let target_provider = self.target_provider(&provider, source_replayable).await?; let scan = session_ctx.read_table(target_provider)?; // Wrap column names in double quotes to preserve case (DataFusion lowercases unquoted identifiers) let on_cols = self @@ -1934,13 +2012,14 @@ impl MergeInsertJob { async fn execute_uncommitted_v2( self, provider: Arc, + source_replayable: bool, ) -> Result<( Transaction, MergeStats, Option, Option, )> { - let plan = self.create_plan(provider).await?; + let plan = self.create_plan(provider, source_replayable).await?; // Execute the plan // Assert that we have exactly one partition since we're designed for single-partition execution @@ -2126,13 +2205,15 @@ impl MergeInsertJob { async fn execute_uncommitted_impl( self, provider: Arc, + source_replayable: bool, ) -> Result { // Check if we can use the fast path let can_use_fast_path = self.can_use_create_plan(provider.schema().as_ref()).await?; if can_use_fast_path { - let (transaction, stats, affected_rows, inserted_rows_filter) = - self.execute_uncommitted_v2(provider).await?; + let (transaction, stats, affected_rows, inserted_rows_filter) = self + .execute_uncommitted_v2(provider, source_replayable) + .await?; return Ok(UncommittedMergeInsert { transaction, affected_rows, @@ -2461,18 +2542,16 @@ impl MergeInsertJob { return Err(Error::not_supported_source("This merge insert configuration does not support explain_plan. Only full-schema merge insert operations without a scalar-index execution path are currently supported.".into())); } - // Create an empty batch with the provided schema to pass to create_plan + // An empty batch with the provided schema stands in for the source. It is + // wrapped in a re-scannable provider so the plan shown is the one the + // execution entry points build; a one-shot source would rule out the + // index-probe target scan. let empty_batch = RecordBatch::new_empty(Arc::new(schema.clone())); - let stream = RecordBatchStreamAdapter::new( - Arc::new(schema.clone()), - futures::stream::once(async { Ok(empty_batch) }).boxed(), - ); + let provider = self.batches_to_provider(vec![empty_batch])?; // Clone self since create_plan consumes the job let cloned_job = self.clone(); - let plan = cloned_job - .create_plan(one_shot_provider(Box::pin(stream))?) - .await?; + let plan = cloned_job.create_plan(provider, true).await?; let display = DisplayableExecutionPlan::new(plan.as_ref()); Ok(format!("{}", display.indent(verbose))) @@ -2503,9 +2582,13 @@ impl MergeInsertJob { return Err(Error::not_supported_source("This merge insert configuration does not support analyze_plan. Only full-schema merge insert operations without a scalar-index execution path are currently supported.".into())); } + // Build the provider the same way execution does, so the plan analyzed is + // the plan that would have run. + let (provider, replayable) = self.stream_source_to_provider(source).await?; + // Clone self since create_plan consumes the job let cloned_job = self.clone(); - let plan = cloned_job.create_plan(one_shot_provider(source)?).await?; + let plan = cloned_job.create_plan(provider, replayable).await?; // Use the analyze_plan function from lance_datafusion, but strip out the wrapper lines let options = LanceExecutionOptions::default(); @@ -2560,6 +2643,9 @@ pub struct UncommittedMergeInsert { struct MergeInsertJobWithProvider { job: MergeInsertJob, provider: Arc, + /// Whether `provider` can be scanned more than once. False only for a + /// one-shot stream that was not spilled, which also disables retries. + replayable: bool, attempt_count: Arc, } @@ -2574,7 +2660,7 @@ impl RetryExecutor for MergeInsertJobWithProvider { // Re-scan the provider on each retry attempt. self.job .clone() - .execute_uncommitted_impl(self.provider.clone()) + .execute_uncommitted_impl(self.provider.clone(), self.replayable) .await } @@ -4738,6 +4824,249 @@ mod tests { assert_eq!(updated_ds.count_rows(None).await.unwrap(), 3); } + /// Dataset with a two-column merge key and a scalar index on `a` only, so + /// the key is partially indexed and the merge takes the v2 plan path. + async fn partially_indexed_dataset(index_on_a: bool) -> (Arc, Arc) { + let initial = record_batch!( + ("a", Int32, [1, 1, 2, 2]), + ("b", Int32, [10, 20, 10, 20]), + ("value", Int32, [100, 200, 300, 400]) + ) + .unwrap(); + let schema = initial.schema(); + let mut ds = Dataset::write( + RecordBatchIterator::new(vec![Ok(initial.clone())], schema.clone()), + "memory://", + None, + ) + .await + .unwrap(); + if index_on_a { + ds.create_index( + &["a"], + IndexType::Scalar, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + } + (Arc::new(ds), schema) + } + + /// Render the merge's physical plan for an empty source of `schema`, stating + /// whether the source provider can be scanned more than once. + async fn merge_plan_string( + job: &MergeInsertJob, + schema: &Arc, + source_replayable: bool, + ) -> String { + let empty = RecordBatch::new_empty(schema.clone()); + let provider: Arc = if source_replayable { + Arc::new(MemTable::try_new(schema.clone(), vec![vec![empty]]).unwrap()) + } else { + let stream = RecordBatchStreamAdapter::new( + schema.clone(), + futures::stream::once(async { Ok(empty) }), + ); + one_shot_provider(Box::pin(stream)).unwrap() + }; + let plan = job + .clone() + .create_plan(provider, source_replayable) + .await + .unwrap(); + format!( + "{}", + DisplayableExecutionPlan::new(plan.as_ref()).indent(true) + ) + } + + /// The v2 target scan reaches matched rows through a scalar index only when + /// doing so cannot miss a row the merge needs. `IndexedLookup` is the probe + /// node; its absence means the plan scans the whole target instead. + #[rstest::rstest] + #[case::partially_indexed(true, true, WhenNotMatchedBySource::Keep, true)] + #[case::use_index_disabled(true, false, WhenNotMatchedBySource::Keep, false)] + #[case::no_index(false, true, WhenNotMatchedBySource::Keep, false)] + #[case::delete_by_source(true, true, WhenNotMatchedBySource::Delete, false)] + #[tokio::test] + async fn test_probe_plan_shape( + #[case] index_on_a: bool, + #[case] use_index: bool, + #[case] by_source: WhenNotMatchedBySource, + #[case] expect_probe: bool, + ) { + let (ds, schema) = partially_indexed_dataset(index_on_a).await; + let job = MergeInsertBuilder::try_new(ds, vec!["a".to_string(), "b".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .when_not_matched_by_source(by_source) + .use_index(use_index) + .try_build() + .unwrap(); + + let plan = merge_plan_string(&job, &schema, true).await; + assert_eq!( + plan.contains("IndexedLookup"), + expect_probe, + "unexpected target scan shape:\n{plan}" + ); + } + + /// The probe scans the source a second time to collect key values, so a + /// source that can only be read once must keep the full-scan hash join. + #[tokio::test] + async fn test_probe_requires_replayable_source() { + let (ds, schema) = partially_indexed_dataset(true).await; + let job = MergeInsertBuilder::try_new(ds, vec!["a".to_string(), "b".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .try_build() + .unwrap(); + + let one_shot = merge_plan_string(&job, &schema, false).await; + assert!( + !one_shot.contains("IndexedLookup"), + "a one-shot source cannot feed the probe:\n{one_shot}" + ); + + let replayable = merge_plan_string(&job, &schema, true).await; + assert!( + replayable.contains("IndexedLookup"), + "a re-scannable source should reach the probe:\n{replayable}" + ); + } + + /// A full-schema upsert produces its output rows entirely from the source, so + /// the probe's take must fetch only the join keys and the row address it needs + /// to supersede the old rows. Pinned because widening this read turns a + /// cheap keyed take into a full-row one. + #[tokio::test] + async fn test_probe_take_reads_only_join_keys() { + let (ds, schema) = partially_indexed_dataset(true).await; + let job = MergeInsertBuilder::try_new(ds, vec!["a".to_string(), "b".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .try_build() + .unwrap(); + + let plan = merge_plan_string(&job, &schema, true).await; + let read = plan + .lines() + .find(|line| line.trim_start().starts_with("LanceRead:")) + .unwrap_or_else(|| panic!("no target read in plan:\n{plan}")); + assert!( + read.contains("projection=[a, b, _rowaddr]"), + "target read should fetch only the keys and row address, got: {read}" + ); + } + + /// Rows in fragments the chosen index does not cover are invisible to the + /// probe, so the plan unions in a scan of those fragments. Without it, updates + /// to rows appended after the index was built are silently dropped. + #[tokio::test] + async fn test_probe_unions_unindexed_fragments() { + let (ds, schema) = partially_indexed_dataset(true).await; + let mut ds = Arc::unwrap_or_clone(ds); + let appended = record_batch!( + ("a", Int32, [3]), + ("b", Int32, [30]), + ("value", Int32, [300]) + ) + .unwrap(); + ds.append( + RecordBatchIterator::new(vec![Ok(appended.clone())], appended.schema()), + None, + ) + .await + .unwrap(); + + let job = MergeInsertBuilder::try_new(Arc::new(ds), vec!["a".to_string(), "b".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .try_build() + .unwrap(); + + let plan = merge_plan_string(&job, &schema, true).await; + assert!( + plan.contains("UnionExec"), + "unindexed fragments must be unioned into the target scan:\n{plan}" + ); + + // One row in an indexed fragment, one in the appended fragment. + let source = record_batch!( + ("a", Int32, [1, 3]), + ("b", Int32, [10, 30]), + ("value", Int32, [999, 333]) + ) + .unwrap(); + let (updated, stats) = job.execute_batches(vec![source]).await.unwrap(); + + assert_eq!( + stats.num_updated_rows, 2, + "the row in the unindexed fragment must also be updated" + ); + assert_eq!(stats.num_inserted_rows, 0); + assert_eq!(updated.count_rows(None).await.unwrap(), 5); + assert_eq!( + updated + .count_rows(Some("a = 3 AND b = 30 AND value = 333".to_string())) + .await + .unwrap(), + 1 + ); + } + + /// The probe emits one candidate set per source batch and does not + /// de-duplicate across batches, so a target row matched by keys in two + /// different batches reaches the sink twice. Source de-duplication is what + /// keeps that from producing a duplicate row address. + #[tokio::test] + async fn test_probe_source_duplicates_across_batches() { + let (ds, _) = partially_indexed_dataset(true).await; + let job = MergeInsertBuilder::try_new(ds, vec!["a".to_string(), "b".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .source_dedupe_behavior(SourceDedupeBehavior::FirstSeen) + .try_build() + .unwrap(); + + // (1, 10) twice in separate batches, so the two probes both match it. + let first = record_batch!( + ("a", Int32, [1]), + ("b", Int32, [10]), + ("value", Int32, [999]) + ) + .unwrap(); + let second = record_batch!( + ("a", Int32, [1]), + ("b", Int32, [10]), + ("value", Int32, [888]) + ) + .unwrap(); + + let (updated, stats) = job.execute_batches(vec![first, second]).await.unwrap(); + + assert_eq!(stats.num_updated_rows, 1); + assert_eq!(stats.num_skipped_duplicates, 1); + assert_eq!(updated.count_rows(None).await.unwrap(), 4); + assert_eq!( + updated + .count_rows(Some("a = 1 AND b = 10 AND value = 999".to_string())) + .await + .unwrap(), + 1, + "first seen source row wins" + ); + } + /// Composite-key delete-only merge_insert (`when_matched(Delete)`, /// `when_not_matched_by_source(Keep)`) removes the matched rows for every /// combination of which join columns carry a scalar index, including when @@ -6720,7 +7049,7 @@ mod tests { let new_data_stream = reader_to_stream(Box::new(new_data)); let plan = merge_insert_job - .create_plan(one_shot_provider(new_data_stream).unwrap()) + .create_plan(one_shot_provider(new_data_stream).unwrap(), false) .await .unwrap(); @@ -6776,7 +7105,7 @@ mod tests { // This should use the fast path (execute_uncommitted_v2) let plan = merge_insert_job - .create_plan(one_shot_provider(new_data_stream).unwrap()) + .create_plan(one_shot_provider(new_data_stream).unwrap(), false) .await .unwrap(); @@ -6827,7 +7156,7 @@ mod tests { let new_data_stream = reader_to_stream(Box::new(new_data_reader)); let plan = merge_insert_job - .create_plan(one_shot_provider(new_data_stream).unwrap()) + .create_plan(one_shot_provider(new_data_stream).unwrap(), false) .await .unwrap(); @@ -6882,7 +7211,7 @@ mod tests { // Should reach the v2 fast path (`create_plan` + FullSchemaMergeInsertExec). // Dropping to v1 here would return an error from create_plan instead. let plan = merge_insert_job - .create_plan(one_shot_provider(new_data_stream).unwrap()) + .create_plan(one_shot_provider(new_data_stream).unwrap(), false) .await .unwrap(); @@ -8013,13 +8342,16 @@ mod tests { // Test explain_plan with default schema (None) let plan = merge_insert_job.explain_plan(None, false).await.unwrap(); - // Also validate the full string structure with pattern matching + // Also validate the full string structure with pattern matching. + // `explain_plan` stands the source up as a materialized table, so its + // exact row count reaches the optimizer and the join collects the + // (smaller) source as its build side — hence source before target. let expected_pattern = "\ MergeInsert: on=[id], when_matched=UpdateAll, when_not_matched=InsertAll, when_not_matched_by_source=Keep... CoalescePartitionsExec... - HashJoinExec... - LanceRead... - StreamingTableExec: partition_sizes=1, projection=[id, name]"; + HashJoinExec: mode=CollectLeft, join_type=Left... + DataSourceExec: partitions=1, partition_sizes=[1] + LanceRead: uri=data, projection=[id], num_fragments=1, range_before=None, range_after=None, row_id=true, row_addr=true, full_filter=--, refine_filter=--"; assert_string_matches(&plan, expected_pattern).unwrap(); // Test with explicit schema @@ -8059,12 +8391,15 @@ MergeInsert: on=[id], when_matched=UpdateAll, when_not_matched=InsertAll, when_n let plan = merge_insert_job.explain_plan(None, false).await.unwrap(); + // `join_type=Left` with the source on the build side is the swapped form + // of `Right` with the target there: either way every source row survives + // the join so unmatched ones can be inserted. let expected_pattern = "\ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_not_matched_by_source=Keep... CoalescePartitionsExec... - HashJoinExec...join_type=Right... - LanceRead... - StreamingTableExec: partition_sizes=1, projection=[id, name]"; + HashJoinExec: mode=CollectLeft, join_type=Left... + DataSourceExec: partitions=1, partition_sizes=[1] + LanceRead: uri=data, projection=[id], num_fragments=1, range_before=None, range_after=None, row_id=true, row_addr=true, full_filter=--, refine_filter=--"; assert_string_matches(&plan, expected_pattern).unwrap(); } @@ -9820,7 +10155,7 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n schema.clone(), ))); let plan = plan_job - .create_plan(one_shot_provider(plan_stream).unwrap()) + .create_plan(one_shot_provider(plan_stream).unwrap(), false) .await .unwrap(); assert_plan_node_equals( @@ -9905,7 +10240,7 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n id_only_schema.clone(), ))); let plan = plan_job - .create_plan(one_shot_provider(plan_stream).unwrap()) + .create_plan(one_shot_provider(plan_stream).unwrap(), false) .await .unwrap(); assert_plan_node_equals( @@ -9990,7 +10325,7 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n schema.clone(), ))); let plan = plan_job - .create_plan(one_shot_provider(plan_stream).unwrap()) + .create_plan(one_shot_provider(plan_stream).unwrap(), false) .await .unwrap(); assert_plan_node_equals( @@ -10098,7 +10433,7 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n schema.clone(), ))); let plan = plan_job - .create_plan(one_shot_provider(plan_stream).unwrap()) + .create_plan(one_shot_provider(plan_stream).unwrap(), false) .await .unwrap(); assert_plan_node_equals( @@ -10221,7 +10556,7 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n schema.clone(), ))); let plan = job - .create_plan(one_shot_provider(plan_stream).unwrap()) + .create_plan(one_shot_provider(plan_stream).unwrap(), false) .await .unwrap(); @@ -11870,7 +12205,7 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n .unwrap(); // The provider's exact row count reaches the plan's statistics. - let plan = job.create_plan(provider).await.unwrap(); + let plan = job.create_plan(provider, true).await.unwrap(); let mut row_counts = Vec::new(); collect_exact_row_counts(&plan, &mut row_counts); assert!( diff --git a/rust/lance/src/dataset/write/merge_insert/probe.rs b/rust/lance/src/dataset/write/merge_insert/probe.rs new file mode 100644 index 00000000000..64d0122896f --- /dev/null +++ b/rust/lance/src/dataset/write/merge_insert/probe.rs @@ -0,0 +1,217 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Scalar-index probe for the merge-insert target table. + +use std::sync::Arc; + +use arrow_schema::{Schema as ArrowSchema, SchemaRef}; +use async_trait::async_trait; +use datafusion::{ + catalog::{Session, TableProvider}, + error::DataFusionError, + logical_expr::{Expr, TableType}, + physical_plan::{ExecutionPlan, coalesce_partitions::CoalescePartitionsExec, union::UnionExec}, +}; +use lance_arrow::SchemaExt; +use lance_core::{ + ROW_ADDR_FIELD, ROW_ID_FIELD, + datatypes::{BlobHandling, OnMissing}, +}; +use lance_table::format::Fragment; + +use crate::Dataset; +use crate::io::exec::{ + filtered_read::{FilteredReadExec, FilteredReadOptions}, + project, + scalar_index::{IndexLookup, MapIndexExec}, +}; + +/// A [`TableProvider`] for the merge-insert target that reaches candidate rows +/// through scalar indices instead of scanning the whole table. +/// +/// It reports the same schema as +/// [`LanceTableProvider`](crate::datafusion::LanceTableProvider) built with row +/// id and row address, so the rest of the merge-insert plan — the join, the +/// action expression, the write sink — is identical either way. Only the target +/// scan differs: +/// +/// ```text +/// Union +/// ├─ FilteredReadExec take the projected columns of probed rows +/// │ └─ CoalescePartitionsExec +/// │ └─ MapIndexExec AND of one IsIn probe per indexed key +/// │ └─ key columns only +/// └─ FilteredReadExec fragments no index covers +/// ``` +/// +/// The probe is allowed to over-match and must not under-match, because the +/// downstream join is what applies the real key predicate. Two things follow +/// from that: +/// +/// - Only the indexed subset of the merge keys needs to be probed. A composite +/// key with one indexed column still prunes by that column, and per-column +/// `IsIn` lists do not correlate values across the tuple anyway. +/// - Fragments that any chosen index does not cover would be invisible to the +/// probe, so they are scanned and unioned in. +/// +/// The source is scanned a second time here to collect key values. Callers must +/// only build this over a re-scannable source. +#[derive(Debug)] +pub(super) struct IndexProbeTarget { + dataset: Arc, + source: Arc, + lookups: Vec, + unindexed_fragments: Arc>, + blob_handling: BlobHandling, + full_schema: SchemaRef, +} + +impl IndexProbeTarget { + pub(super) fn try_new( + dataset: Arc, + source: Arc, + lookups: Vec, + unindexed_fragments: Vec, + blob_handling: BlobHandling, + ) -> crate::Result { + if lookups.is_empty() { + return Err(crate::Error::internal( + "IndexProbeTarget requires at least one index lookup", + )); + } + let full_schema = ArrowSchema::from(dataset.schema()) + .try_with_column(ROW_ID_FIELD.clone())? + .try_with_column(ROW_ADDR_FIELD.clone())?; + Ok(Self { + dataset, + source, + lookups, + unindexed_fragments: Arc::new(unindexed_fragments), + blob_handling, + full_schema: Arc::new(full_schema), + }) + } + + /// The columns DataFusion asked for, in the order it expects them back. + fn projected_schema( + &self, + projection: Option<&Vec>, + ) -> datafusion::common::Result { + match projection { + Some(indices) => Ok(self.full_schema.project(indices)?), + None => Ok(self.full_schema.as_ref().clone()), + } + } + + /// Scan the source's key columns and turn them into candidate row addresses. + async fn probe( + &self, + state: &dyn Session, + columns: &lance_core::datatypes::Projection, + output_schema: &ArrowSchema, + ) -> datafusion::common::Result> { + let source_schema = self.source.schema(); + let mut key_indices = self + .lookups + .iter() + .map(|lookup| source_schema.index_of(&lookup.column)) + .collect::, _>>()?; + key_indices.sort_unstable(); + key_indices.dedup(); + let keys = self + .source + .scan(state, Some(&key_indices), &[], None) + .await?; + + // `MapIndexExec` matches its input columns to `lookups` positionally and + // reads a single partition, so the keys are reordered and coalesced first. + let lookup_schema = ArrowSchema::new( + self.lookups + .iter() + .map(|lookup| source_schema.field_with_name(&lookup.column).cloned()) + .collect::, _>>()?, + ); + let keys = Arc::new(project(keys, &lookup_schema)?); + let keys = Arc::new(CoalescePartitionsExec::new(keys)); + + let probe = Arc::new(MapIndexExec::new_multi( + self.dataset.clone(), + self.lookups.clone(), + keys, + )); + let take = Arc::new(FilteredReadExec::try_new( + self.dataset.clone(), + FilteredReadOptions::new(columns.clone()), + Some(probe), + )?); + Ok(Arc::new(project(take, output_schema)?)) + } + + /// Scan the fragments the probe cannot see. + fn scan_unindexed( + &self, + columns: &lance_core::datatypes::Projection, + output_schema: &ArrowSchema, + ) -> datafusion::common::Result> { + let options = FilteredReadOptions::new(columns.clone()) + .with_fragments(self.unindexed_fragments.clone()); + let scan = Arc::new(FilteredReadExec::try_new( + self.dataset.clone(), + options, + None, + )?); + Ok(Arc::new(project(scan, output_schema)?)) + } +} + +#[async_trait] +impl TableProvider for IndexProbeTarget { + fn schema(&self) -> SchemaRef { + self.full_schema.clone() + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + state: &dyn Session, + projection: Option<&Vec>, + filters: &[Expr], + limit: Option, + ) -> datafusion::common::Result> { + // Neither is reachable from the merge-insert plan, and both would be + // wrong to drop: `supports_filters_pushdown` rejects every filter, and + // nothing puts a limit above the target scan. + if !filters.is_empty() { + return Err(DataFusionError::Internal(format!( + "IndexProbeTarget cannot apply pushed-down filters, got {} of them", + filters.len() + ))); + } + if let Some(limit) = limit { + return Err(DataFusionError::Internal(format!( + "IndexProbeTarget cannot apply a pushed-down limit of {limit}" + ))); + } + + let output_schema = self.projected_schema(projection)?; + let columns = self + .dataset + .empty_projection() + .with_blob_handling(self.blob_handling.clone()) + .union_columns( + output_schema.fields().iter().map(|field| field.name()), + OnMissing::Error, + )?; + + let probe = self.probe(state, &columns, &output_schema).await?; + if self.unindexed_fragments.is_empty() { + return Ok(probe); + } + let unindexed = self.scan_unindexed(&columns, &output_schema)?; + Ok(UnionExec::try_new(vec![probe, unindexed])?) + } +} From 637744987254154db71f1f6c3c62a4aee1d93d8c Mon Sep 17 00:00:00 2001 From: Will Jones Date: Tue, 28 Jul 2026 14:48:15 -0700 Subject: [PATCH 3/3] fix(merge_insert): make the index probe see every source key exactly once Benchmarking the probe surfaced two ways a partially-indexed composite key merge produced wrong results. A materialized source is spread across partitions and `MapIndexExec` probes only the partition it is asked for, so the keys in every other partition were never probed and their target rows were inserted as duplicates instead of updated. `MapIndexExec` now requires a single input partition, which is what keeps the optimizer from removing the probe's `CoalescePartitionsExec`. The probe also evaluates one query per source batch, and an over-matching probe can reach the same target row from two batches. The target row was then read twice, giving one source row two candidate matches, and the merge failed with "Ambiguous merge inserts are prohibited". Candidate row addresses are now de-duplicated before the take. Co-Authored-By: Claude Opus 5 (1M context) --- rust/lance/src/dataset/write/merge_insert.rs | 72 ++++++++-- .../src/dataset/write/merge_insert/probe.rs | 132 +++++++++++++++++- rust/lance/src/io/exec/scalar_index.rs | 14 ++ 3 files changed, 204 insertions(+), 14 deletions(-) diff --git a/rust/lance/src/dataset/write/merge_insert.rs b/rust/lance/src/dataset/write/merge_insert.rs index 57d57eb0d3a..c06f6dd0800 100644 --- a/rust/lance/src/dataset/write/merge_insert.rs +++ b/rust/lance/src/dataset/write/merge_insert.rs @@ -4909,11 +4909,13 @@ mod tests { .unwrap(); let plan = merge_plan_string(&job, &schema, true).await; - assert_eq!( - plan.contains("IndexedLookup"), - expect_probe, - "unexpected target scan shape:\n{plan}" - ); + for node in ["IndexedLookup", "DistinctRowAddrs"] { + assert_eq!( + plan.contains(node), + expect_probe, + "unexpected target scan shape, looking for {node}:\n{plan}" + ); + } } /// The probe scans the source a second time to collect key values, so a @@ -5023,10 +5025,9 @@ mod tests { ); } - /// The probe emits one candidate set per source batch and does not - /// de-duplicate across batches, so a target row matched by keys in two - /// different batches reaches the sink twice. Source de-duplication is what - /// keeps that from producing a duplicate row address. + /// The same key in two source batches makes both probes name the same target + /// row. Source de-duplication settles which source row wins; the probe's own + /// de-duplication is what keeps the target row from being read twice. #[tokio::test] async fn test_probe_source_duplicates_across_batches() { let (ds, _) = partially_indexed_dataset(true).await; @@ -5067,6 +5068,59 @@ mod tests { ); } + /// One source batch per key, which is the shape both probe hazards need. + /// + /// A materialized source is spread across partitions and the probe reads one + /// partition, so the key columns have to be coalesced first — otherwise the + /// keys in every other partition go unprobed and their target rows are + /// inserted as duplicates instead of updated. And because only `a` is + /// indexed, the batches for `(1, 10)` and `(1, 20)` both probe `a IsIn [1]` + /// and name each other's target row, so the candidates have to be + /// de-duplicated or the merge fails as ambiguous. + #[tokio::test] + async fn test_probe_multi_partition_source() { + let (ds, _) = partially_indexed_dataset(true).await; + let job = MergeInsertBuilder::try_new(ds, vec!["a".to_string(), "b".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .try_build() + .unwrap(); + + let keys = [(1, 10), (1, 20), (2, 10), (2, 20)]; + let batches = keys + .iter() + .enumerate() + .map(|(idx, (a, b))| { + record_batch!( + ("a", Int32, [*a]), + ("b", Int32, [*b]), + ("value", Int32, [900 + idx as i32]) + ) + .unwrap() + }) + .collect(); + + let (updated, stats) = job.execute_batches(batches).await.unwrap(); + + assert_eq!(stats.num_updated_rows, keys.len() as u64); + assert_eq!(stats.num_inserted_rows, 0); + assert_eq!(updated.count_rows(None).await.unwrap(), keys.len()); + for (idx, (a, b)) in keys.iter().enumerate() { + assert_eq!( + updated + .count_rows(Some(format!( + "a = {a} AND b = {b} AND value = {}", + 900 + idx + ))) + .await + .unwrap(), + 1, + "row ({a}, {b}) from source batch {idx} was not updated" + ); + } + } + /// Composite-key delete-only merge_insert (`when_matched(Delete)`, /// `when_not_matched_by_source(Keep)`) removes the matched rows for every /// combination of which join columns carry a scalar index, including when diff --git a/rust/lance/src/dataset/write/merge_insert/probe.rs b/rust/lance/src/dataset/write/merge_insert/probe.rs index 64d0122896f..f6193115bbf 100644 --- a/rust/lance/src/dataset/write/merge_insert/probe.rs +++ b/rust/lance/src/dataset/write/merge_insert/probe.rs @@ -5,28 +5,145 @@ use std::sync::Arc; +use arrow_array::{RecordBatch, UInt64Array, cast::AsArray, types::UInt64Type}; use arrow_schema::{Schema as ArrowSchema, SchemaRef}; use async_trait::async_trait; use datafusion::{ catalog::{Session, TableProvider}, error::DataFusionError, + execution::TaskContext, logical_expr::{Expr, TableType}, - physical_plan::{ExecutionPlan, coalesce_partitions::CoalescePartitionsExec, union::UnionExec}, + physical_expr::{Distribution, EquivalenceProperties}, + physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, + SendableRecordBatchStream, + coalesce_partitions::CoalescePartitionsExec, + execution_plan::{Boundedness, EmissionType}, + stream::RecordBatchStreamAdapter, + union::UnionExec, + }, }; +use futures::TryStreamExt; use lance_arrow::SchemaExt; use lance_core::{ ROW_ADDR_FIELD, ROW_ID_FIELD, datatypes::{BlobHandling, OnMissing}, }; use lance_table::format::Fragment; +use roaring::RoaringTreemap; use crate::Dataset; use crate::io::exec::{ filtered_read::{FilteredReadExec, FilteredReadOptions}, project, - scalar_index::{IndexLookup, MapIndexExec}, + scalar_index::{INDEX_LOOKUP_SCHEMA, IndexLookup, MapIndexExec}, }; +/// Drops row addresses an earlier batch already emitted, so a probe's candidate +/// stream is a set. +/// +/// [`MapIndexExec`] evaluates one query per input batch and emits that batch's +/// matches, so an over-matching probe can reach the same target row from two +/// different batches. Reading a target row twice would give one source row two +/// candidate matches in the join, and the merge rejects that as ambiguous — so +/// the duplicates have to go before the take. +#[derive(Debug)] +struct DistinctRowAddrsExec { + input: Arc, + properties: Arc, +} + +impl DistinctRowAddrsExec { + fn new(input: Arc) -> Self { + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(INDEX_LOOKUP_SCHEMA.clone()), + Partitioning::RoundRobinBatch(1), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { input, properties } + } + + fn retain_unseen(seen: &mut RoaringTreemap, batch: &RecordBatch) -> RecordBatch { + let addrs = batch.column(0).as_primitive::(); + let unseen: UInt64Array = addrs + .values() + .iter() + .copied() + .filter(|addr| seen.insert(*addr)) + .collect(); + RecordBatch::try_new(INDEX_LOOKUP_SCHEMA.clone(), vec![Arc::new(unseen)]) + .expect("a UInt64 column always matches INDEX_LOOKUP_SCHEMA") + } +} + +impl DisplayAs for DistinctRowAddrsExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "DistinctRowAddrs") + } + DisplayFormatType::TreeRender => write!(f, "DistinctRowAddrs"), + } + } +} + +impl ExecutionPlan for DistinctRowAddrsExec { + fn name(&self) -> &str { + "DistinctRowAddrsExec" + } + + fn schema(&self) -> SchemaRef { + INDEX_LOOKUP_SCHEMA.clone() + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> datafusion::common::Result> { + let [input] = <[_; 1]>::try_from(children).map_err(|children| { + DataFusionError::Internal(format!( + "DistinctRowAddrsExec requires exactly one child, got {}", + children.len() + )) + })?; + Ok(Arc::new(Self::new(input))) + } + + fn required_input_distribution(&self) -> Vec { + // De-duplicating across batches only works if every batch arrives here. + vec![Distribution::SinglePartition] + } + + fn benefits_from_input_partitioning(&self) -> Vec { + vec![false] + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> datafusion::common::Result { + let mut seen = RoaringTreemap::new(); + let stream = self + .input + .execute(partition, context)? + .map_ok(move |batch| Self::retain_unseen(&mut seen, &batch)); + Ok(Box::pin(RecordBatchStreamAdapter::new( + INDEX_LOOKUP_SCHEMA.clone(), + stream, + ))) + } + + fn properties(&self) -> &Arc { + &self.properties + } +} + /// A [`TableProvider`] for the merge-insert target that reaches candidate rows /// through scalar indices instead of scanning the whole table. /// @@ -39,14 +156,15 @@ use crate::io::exec::{ /// ```text /// Union /// ├─ FilteredReadExec take the projected columns of probed rows -/// │ └─ CoalescePartitionsExec +/// │ └─ DistinctRowAddrsExec one candidate row address at most once /// │ └─ MapIndexExec AND of one IsIn probe per indexed key -/// │ └─ key columns only +/// │ └─ CoalescePartitionsExec +/// │ └─ key columns only /// └─ FilteredReadExec fragments no index covers /// ``` /// /// The probe is allowed to over-match and must not under-match, because the -/// downstream join is what applies the real key predicate. Two things follow +/// downstream join is what applies the real key predicate. Three things follow /// from that: /// /// - Only the indexed subset of the merge keys needs to be probed. A composite @@ -54,6 +172,9 @@ use crate::io::exec::{ /// `IsIn` lists do not correlate values across the tuple anyway. /// - Fragments that any chosen index does not cover would be invisible to the /// probe, so they are scanned and unioned in. +/// - Over-matching is only harmless while each candidate row reaches the join +/// once. Probes are evaluated per source batch, so two batches can name the +/// same target row and the candidates are de-duplicated before the take. /// /// The source is scanned a second time here to collect key values. Callers must /// only build this over a re-scannable source. @@ -140,6 +261,7 @@ impl IndexProbeTarget { self.lookups.clone(), keys, )); + let probe = Arc::new(DistinctRowAddrsExec::new(probe)); let take = Arc::new(FilteredReadExec::try_new( self.dataset.clone(), FilteredReadOptions::new(columns.clone()), diff --git a/rust/lance/src/io/exec/scalar_index.rs b/rust/lance/src/io/exec/scalar_index.rs index b943ad12d3a..97891feeb68 100644 --- a/rust/lance/src/io/exec/scalar_index.rs +++ b/rust/lance/src/io/exec/scalar_index.rs @@ -17,6 +17,7 @@ use arrow_schema::{Schema, SchemaRef}; use async_recursion::async_recursion; use async_trait::async_trait; use datafusion::{ + physical_expr::Distribution, physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, execution_plan::{Boundedness, EmissionType}, @@ -583,6 +584,19 @@ impl ExecutionPlan for MapIndexExec { fn supports_limit_pushdown(&self) -> bool { false } + + fn required_input_distribution(&self) -> Vec { + // `execute` probes exactly the partition it is asked for, and this node + // declares a single output partition, so every input batch has to arrive + // on partition 0. Without this the optimizer is free to leave a + // multi-partition input in place and the keys in every partition but the + // first are never probed. + vec![Distribution::SinglePartition] + } + + fn benefits_from_input_partitioning(&self) -> Vec { + vec![false] + } } pub static MATERIALIZE_INDEX_SCHEMA: LazyLock =