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() 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..c06f6dd0800 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,303 @@ 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; + 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 + /// 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 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; + 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" + ); + } + + /// 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 @@ -6720,7 +7103,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 +7159,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 +7210,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 +7265,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 +8396,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 +8445,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 +10209,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 +10294,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 +10379,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 +10487,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 +10610,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 +12259,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..f6193115bbf --- /dev/null +++ b/rust/lance/src/dataset/write/merge_insert/probe.rs @@ -0,0 +1,339 @@ +// 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_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_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::{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. +/// +/// 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 +/// │ └─ DistinctRowAddrsExec one candidate row address at most once +/// │ └─ MapIndexExec AND of one IsIn probe per indexed key +/// │ └─ 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. Three 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. +/// - 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. +#[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 probe = Arc::new(DistinctRowAddrsExec::new(probe)); + 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])?) + } +} 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 =