mirror of
https://github.com/scylladb/scylladb.git
synced 2026-05-13 03:12:13 +00:00
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:
@@ -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
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user