From 4ba24bf64b30ef9f8ab258a7030da4bf921c6c61 Mon Sep 17 00:00:00 2001 From: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Sat, 8 Aug 2026 12:35:32 +0000 Subject: [PATCH] test(rust): detect indexed delete compilation regressions --- rust/lancedb/src/table/delete.rs | 45 +------ rust/lancedb/tests/indexed_delete_test.rs | 152 ++++++++++++++++++++++ 2 files changed, 154 insertions(+), 43 deletions(-) create mode 100644 rust/lancedb/tests/indexed_delete_test.rs diff --git a/rust/lancedb/src/table/delete.rs b/rust/lancedb/src/table/delete.rs index d57ae9cf7..8f11ee019 100644 --- a/rust/lancedb/src/table/delete.rs +++ b/rust/lancedb/src/table/delete.rs @@ -62,12 +62,11 @@ pub(crate) async fn execute_delete( #[cfg(test)] mod tests { use crate::connect; - use arrow_array::{Int32Array, RecordBatch, StringArray, record_batch}; + use arrow_array::{Int32Array, RecordBatch, record_batch}; use arrow_schema::{DataType, Field, Schema}; use std::sync::Arc; - use crate::index::Index; - use crate::query::{ExecutableQuery, QueryBase}; + use crate::query::ExecutableQuery; use futures::TryStreamExt; #[tokio::test] async fn test_delete_simple() { @@ -168,46 +167,6 @@ mod tests { assert_eq!(table.count_rows(None).await.unwrap(), 0); } - #[tokio::test] - async fn test_delete_large_in_list_with_btree_index() { - 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..10_000).map(|id| format!("id_{id}"))); - let batch = RecordBatch::try_new(schema, vec![Arc::new(ids)]).unwrap(); - let table = conn - .create_table("test_delete_large_in_list", batch) - .execute() - .await - .unwrap(); - - table - .create_index(&["id"], Index::BTree(Default::default())) - .execute() - .await - .unwrap(); - - let values = (0..10_000) - .step_by(10) - .map(|id| format!("'id_{id}'")) - .collect::>() - .join(","); - let predicate = format!("id IN ({values})"); - - // A large IN-list must stay on the scalar-index path. Lance compiles this - // predicate once and reuses it across all BTree pages. - let plan = table - .query() - .only_if(&predicate) - .explain_plan(false) - .await - .unwrap(); - assert!(plan.contains("ScalarIndexQuery"), "unexpected plan: {plan}"); - - let result = table.delete(&predicate).await.unwrap(); - assert_eq!(result.num_deleted_rows, 1_000); - assert_eq!(table.count_rows(None).await.unwrap(), 9_000); - } - #[tokio::test] async fn test_delete_false_increments_version() { let conn = connect("memory://").execute().await.unwrap(); 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" + ); +}