From f74a7d38cae4ee92b00e80f720f3e90e388931f4 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 16:32:37 +0800 Subject: [PATCH] test(rust): detect indexed delete compilation regressions ## Summary - Compare equal-size 1,000-value deletes confined to one BTree page and spread across many pages. - Use an isolated current-thread allocation counter instead of a wall-clock threshold. - Warm index pages first so the assertion detects repeated predicate materialization rather than I/O. ## Root cause Lance previously rebuilt a large `IN` predicate for every BTree page, making indexed lookups scale as `O(pages x values)`. The pinned Lance revision compiles the predicate once and reuses it across pages. The regression guard observes the allocation scaling caused by that invariant directly. Fixes #2085 --- rust/lancedb/tests/indexed_delete_test.rs | 152 ++++++++++++++++++++++ 1 file changed, 152 insertions(+) create mode 100644 rust/lancedb/tests/indexed_delete_test.rs diff --git a/rust/lancedb/tests/indexed_delete_test.rs b/rust/lancedb/tests/indexed_delete_test.rs new file mode 100644 index 000000000..4ebaf5238 --- /dev/null +++ b/rust/lancedb/tests/indexed_delete_test.rs @@ -0,0 +1,152 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use std::{ + alloc::{GlobalAlloc, Layout, System}, + cell::Cell, + future::Future, + sync::Arc, +}; + +use arrow_array::{RecordBatch, StringArray}; +use arrow_schema::{DataType, Field, Schema}; +use futures::TryStreamExt; +use lancedb::{ + Table, connect, + index::Index, + query::{ExecutableQuery, QueryBase}, +}; + +struct ThreadCountingAllocator; + +thread_local! { + static COUNT_ALLOCATIONS: Cell = const { Cell::new(false) }; + static ALLOCATED_BYTES: Cell = const { Cell::new(0) }; +} + +unsafe impl GlobalAlloc for ThreadCountingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let ptr = unsafe { System.alloc(layout) }; + if !ptr.is_null() { + record_allocation(layout.size()); + } + ptr + } + + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + let ptr = unsafe { System.alloc_zeroed(layout) }; + if !ptr.is_null() { + record_allocation(layout.size()); + } + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + unsafe { System.dealloc(ptr, layout) }; + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + let new_ptr = unsafe { System.realloc(ptr, layout, new_size) }; + if !new_ptr.is_null() { + record_allocation(new_size); + } + new_ptr + } +} + +#[global_allocator] +static ALLOCATOR: ThreadCountingAllocator = ThreadCountingAllocator; + +const ROW_COUNT: usize = 262_144; +const VALUE_COUNT: usize = 1_000; + +fn record_allocation(bytes: usize) { + COUNT_ALLOCATIONS.with(|enabled| { + if enabled.get() { + ALLOCATED_BYTES.with(|allocated| allocated.set(allocated.get() + bytes)); + } + }); +} + +async fn measure_allocated_bytes(future: F) -> (F::Output, usize) { + ALLOCATED_BYTES.with(|allocated| allocated.set(0)); + COUNT_ALLOCATIONS.with(|enabled| enabled.set(true)); + let output = future.await; + COUNT_ALLOCATIONS.with(|enabled| enabled.set(false)); + let allocated = ALLOCATED_BYTES.with(Cell::get); + (output, allocated) +} + +fn in_predicate(ids: impl Iterator) -> String { + let values = ids + .map(|id| format!("'id_{id:06}'")) + .collect::>() + .join(","); + format!("id IN ({values})") +} + +async fn create_indexed_table(name: &str) -> Table { + let conn = connect("memory://").execute().await.unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Utf8, false)])); + let ids = StringArray::from_iter_values((0..ROW_COUNT).map(|id| format!("id_{id:06}"))); + let batch = RecordBatch::try_new(schema, vec![Arc::new(ids)]).unwrap(); + let table = conn.create_table(name, batch).execute().await.unwrap(); + table + .create_index(&["id"], Index::BTree(Default::default())) + .execute() + .await + .unwrap(); + table +} + +async fn warm_index(table: &Table, predicate: &str) { + table + .query() + .only_if(predicate) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); +} + +#[tokio::test(flavor = "current_thread")] +async fn large_in_delete_compiles_predicate_once() { + let clustered = in_predicate(0..VALUE_COUNT); + let spread = in_predicate((0..VALUE_COUNT).map(|id| id * (ROW_COUNT / VALUE_COUNT))); + let clustered_table = create_indexed_table("clustered_ids").await; + let spread_table = create_indexed_table("spread_ids").await; + + let plan = spread_table + .query() + .only_if(&spread) + .explain_plan(false) + .await + .unwrap(); + assert!(plan.contains("ScalarIndexQuery"), "unexpected plan: {plan}"); + + // Remove page-loading noise from the allocation comparison. Predicate + // compilation is deliberately not cached, so each delete still compiles it. + warm_index(&clustered_table, &spread).await; + warm_index(&spread_table, &spread).await; + + let (clustered_result, clustered_bytes) = + measure_allocated_bytes(clustered_table.delete(&clustered)).await; + let (spread_result, spread_bytes) = measure_allocated_bytes(spread_table.delete(&spread)).await; + + assert_eq!( + clustered_result.unwrap().num_deleted_rows, + VALUE_COUNT as u64 + ); + assert_eq!(spread_result.unwrap().num_deleted_rows, VALUE_COUNT as u64); + + // Both predicates contain the same number and size of values. Spreading them + // across BTree pages may add modest page-processing overhead, but it must not + // rematerialize all values per page. This ratio fails by a wide margin if + // Lance's compile-once path is moved back inside the per-page loop. + assert!( + spread_bytes * 2 < clustered_bytes * 3, + "spread delete allocated {spread_bytes} bytes versus {clustered_bytes} for one page" + ); +}