vector_search: test: migrate paging warnings tests to Python

Move the paging warning related tests from C++ vector_store_client_test to
Python test_vector_search_with_vector_store_mock.
This commit is contained in:
Karol Nowacki
2026-04-27 09:37:08 +00:00
parent 84787ce6a5
commit 20b953ef8c
2 changed files with 49 additions and 76 deletions

View File

@@ -19,6 +19,7 @@ import threading
import pytest
from cassandra.protocol import InvalidRequest
from cassandra.query import SimpleStatement
from test.pylib.skip_types import skip_env
from .util import config_value_context, local_process_id, new_test_table, unique_name, is_scylla
@@ -196,3 +197,51 @@ def test_vector_search_local_vector_index_create_and_query_do_not_fail(cql, test
cql.execute(
f"SELECT * FROM {table} WHERE pk1 = 1 AND pk2 = 2 ORDER BY embedding ANN OF [0.1, 0.2, 0.3] LIMIT 5")
# Verify that a paging warning is emitted when page size is smaller than LIMIT.
def test_vector_search_paging_warning_when_page_size_smaller_than_limit(cql, test_keyspace, vector_store_mock, skip_without_tablets):
schema = "pk1 tinyint, pk2 tinyint, ck1 tinyint, ck2 tinyint, embedding vector<float, 3>, PRIMARY KEY ((pk1, pk2), ck1, ck2)"
with new_test_table(cql, test_keyspace, schema) as table:
cql.execute(
f"CREATE CUSTOM INDEX ON {table}(embedding) USING 'vector_index'")
result = cql.execute(SimpleStatement(
f"SELECT * FROM {table} ORDER BY embedding ANN OF [0.1, 0.2, 0.3] LIMIT 100", fetch_size=5))
warnings = result.response_future.warnings
assert warnings
assert len(warnings) == 1
assert "Paging is not supported for Vector Search queries. The entire result set has been returned." == warnings[0]
# Verify no paging warning is emitted when paging is disabled (fetch_size=0).
def test_vector_search_no_paging_warning_when_paging_disabled(cql, test_keyspace, vector_store_mock, skip_without_tablets):
schema = "pk1 tinyint, pk2 tinyint, ck1 tinyint, ck2 tinyint, embedding vector<float, 3>, PRIMARY KEY ((pk1, pk2), ck1, ck2)"
with new_test_table(cql, test_keyspace, schema) as table:
cql.execute(
f"CREATE CUSTOM INDEX ON {table}(embedding) USING 'vector_index'")
result = cql.execute(SimpleStatement(
f"SELECT * FROM {table} ORDER BY embedding ANN OF [0.1, 0.2, 0.3] LIMIT 100", fetch_size=0))
assert not result.response_future.warnings
# Verify no paging warning is emitted when LIMIT is less than page size.
def test_vector_search_no_paging_warning_when_limit_less_than_page_size(cql, test_keyspace, vector_store_mock, skip_without_tablets):
schema = "pk1 tinyint, pk2 tinyint, ck1 tinyint, ck2 tinyint, embedding vector<float, 3>, PRIMARY KEY ((pk1, pk2), ck1, ck2)"
with new_test_table(cql, test_keyspace, schema) as table:
cql.execute(
f"CREATE CUSTOM INDEX ON {table}(embedding) USING 'vector_index'")
result = cql.execute(SimpleStatement(
f"SELECT * FROM {table} ORDER BY embedding ANN OF [0.1, 0.2, 0.3] LIMIT 5", fetch_size=100))
assert not result.response_future.warnings

View File

@@ -709,82 +709,6 @@ SEASTAR_TEST_CASE(vector_search_metrics_test) {
cfg);
}
SEASTAR_TEST_CASE(vector_store_client_test_paging_warning) {
auto s1 = co_await make_vs_mock_server();
auto cfg = make_config();
cfg.db_config->vector_store_primary_uri.set(format("http://s1.node:{}", s1->port()));
co_await do_with_cql_env(
[&s1](cql_test_env& env) -> future<> {
auto schema = co_await create_test_table(env, "ks", "test");
auto& vs = env.local_qp().vector_store_client();
configure(vs).with_dns({{"s1.node", std::vector<std::string>{s1->host()}}});
vs.start_background_tasks();
auto result = co_await env.execute_cql("CREATE CUSTOM INDEX idx ON ks.test (embedding) USING 'vector_index'");
auto qo = std::make_unique<cql3::query_options>(db::consistency_level::LOCAL_ONE, std::vector<cql3::raw_value>{},
cql3::query_options::specific_options{5, nullptr, {}, api::new_timestamp()});
auto msg = co_await env.execute_cql("SELECT * FROM ks.test ORDER BY embedding ANN OF [0.1, 0.2, 0.3] LIMIT 100;", std::move(qo));
auto warns = msg->warnings();
BOOST_REQUIRE_EQUAL(warns.size(), 1);
BOOST_CHECK(warns[0] == "Paging is not supported for Vector Search queries. The entire result set has been returned.");
},
cfg)
.finally([&s1] {
return s1->stop();
});
}
SEASTAR_TEST_CASE(vector_store_client_test_paging_warning_doesnt_show_when_paging_disabled) {
auto s1 = co_await make_vs_mock_server();
auto cfg = make_config();
cfg.db_config->vector_store_primary_uri.set(format("http://s1.node:{}", s1->port()));
co_await do_with_cql_env(
[&s1](cql_test_env& env) -> future<> {
auto schema = co_await create_test_table(env, "ks", "test");
auto& vs = env.local_qp().vector_store_client();
configure(vs).with_dns({{"s1.node", std::vector<std::string>{s1->host()}}});
vs.start_background_tasks();
auto result = co_await env.execute_cql("CREATE CUSTOM INDEX idx ON ks.test (embedding) USING 'vector_index'");
auto qo = std::make_unique<cql3::query_options>(db::consistency_level::LOCAL_ONE, std::vector<cql3::raw_value>{},
cql3::query_options::specific_options{0, nullptr, {}, api::new_timestamp()});
auto msg = co_await env.execute_cql("SELECT * FROM ks.test ORDER BY embedding ANN OF [0.1, 0.2, 0.3] LIMIT 100;", std::move(qo));
auto warns = msg->warnings();
BOOST_REQUIRE_EQUAL(warns.size(), 0);
},
cfg)
.finally([&s1] {
return s1->stop();
});
}
SEASTAR_TEST_CASE(vector_store_client_test_paging_warning_doesnt_show_when_limit_less_than_page_size) {
auto s1 = co_await make_vs_mock_server();
auto cfg = make_config();
cfg.db_config->vector_store_primary_uri.set(format("http://s1.node:{}", s1->port()));
co_await do_with_cql_env(
[&s1](cql_test_env& env) -> future<> {
auto schema = co_await create_test_table(env, "ks", "test");
auto& vs = env.local_qp().vector_store_client();
configure(vs).with_dns({{"s1.node", std::vector<std::string>{s1->host()}}});
vs.start_background_tasks();
auto result = co_await env.execute_cql("CREATE CUSTOM INDEX idx ON ks.test (embedding) USING 'vector_index'");
auto qo = std::make_unique<cql3::query_options>(db::consistency_level::LOCAL_ONE, std::vector<cql3::raw_value>{},
cql3::query_options::specific_options{100, nullptr, {}, api::new_timestamp()});
auto msg = co_await env.execute_cql("SELECT * FROM ks.test ORDER BY embedding ANN OF [0.1, 0.2, 0.3] LIMIT 5;", std::move(qo));
auto warns = msg->warnings();
BOOST_REQUIRE_EQUAL(warns.size(), 0);
},
cfg)
.finally([&s1] {
return s1->stop();
});
}
SEASTAR_TEST_CASE(vector_store_client_node_recovery_after_backoff) {
auto unavail_server = co_await make_unavailable_server();
std::unique_ptr<vs_mock_server> avail_server;