diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index 6840be052..6c232f6c3 100644 --- a/python/python/tests/test_query.py +++ b/python/python/tests/test_query.py @@ -779,6 +779,56 @@ async def test_distance_range_with_new_rows_async(): assert dist >= min_dist +def test_flat_vector_search_large_limit_is_globally_sorted_after_upsert(mem_db): + # Regression test for https://github.com/lancedb/lancedb/issues/3669. + # Limits above Lance's 8,192-row output batch size used to expose batches in + # scheduling order instead of global distance order after plan repartitioning. + rng = np.random.default_rng(42) + num_rows = 20_000 + fragment_size = 5_000 + dimension = 16 + limit = 12_000 + vectors = rng.random((num_rows, dimension), dtype=np.float32) + + def make_data(start, end): + return pa.table( + { + "id": pa.array(np.arange(start, end), type=pa.int64()), + "vector": pa.FixedSizeListArray.from_arrays( + pa.array(vectors[start:end].reshape(-1)), dimension + ), + } + ) + + table = mem_db.create_table("flat_knn_order", make_data(0, fragment_size)) + for start in range(fragment_size, num_rows, fragment_size): + table.add(make_data(start, start + fragment_size)) + + upsert = pa.table( + { + "id": pa.array([0], type=pa.int64()), + "vector": pa.array( + [[0.0] * dimension], type=pa.list_(pa.float32(), dimension) + ), + } + ) + ( + table.merge_insert("id") + .when_matched_update_all() + .when_not_matched_insert_all() + .execute(upsert) + ) + + query = np.zeros(dimension, dtype=np.float32) + for _ in range(10): + result = ( + table.search(query).where("id >= 0").select(["id"]).limit(limit).to_arrow() + ) + distances = result["_distance"].to_numpy() + assert len(distances) == limit + assert np.all(distances[:-1] <= distances[1:]) + + @pytest.mark.parametrize( "multivec_table", [pa.float16(), pa.float32(), pa.float64()], indirect=True )