diff --git a/test/cqlpy/test_vector_search_with_vector_store_mock.py b/test/cqlpy/test_vector_search_with_vector_store_mock.py index 03d6a56f19..83a3082ca1 100644 --- a/test/cqlpy/test_vector_search_with_vector_store_mock.py +++ b/test/cqlpy/test_vector_search_with_vector_store_mock.py @@ -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, 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, 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, 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 diff --git a/test/vector_search/vector_store_client_test.cc b/test/vector_search/vector_store_client_test.cc index cbf2c82b91..43a7fcd56c 100644 --- a/test/vector_search/vector_store_client_test.cc +++ b/test/vector_search/vector_store_client_test.cc @@ -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{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(db::consistency_level::LOCAL_ONE, std::vector{}, - 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{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(db::consistency_level::LOCAL_ONE, std::vector{}, - 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{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(db::consistency_level::LOCAL_ONE, std::vector{}, - 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 avail_server;