diff --git a/cql3/restrictions/statement_restrictions.hh b/cql3/restrictions/statement_restrictions.hh index af727f7fb0..b6e74592cb 100644 --- a/cql3/restrictions/statement_restrictions.hh +++ b/cql3/restrictions/statement_restrictions.hh @@ -201,6 +201,10 @@ public: return _clustering_columns_restrictions; } + const expr::expression& get_nonprimary_key_restrictions() const { + return _nonprimary_key_restrictions; + } + // Get a set of columns restricted by the IS NOT NULL restriction. // IS NOT NULL is a special case that is handled separately from other restrictions. const std::unordered_set get_not_null_columns() const; diff --git a/test/vector_search/filter_test.cc b/test/vector_search/filter_test.cc index 6bdc89ab4d..5b572a3028 100644 --- a/test/vector_search/filter_test.cc +++ b/test/vector_search/filter_test.cc @@ -407,4 +407,44 @@ SEASTAR_TEST_CASE(to_json_no_bind_markers_uses_cache) { }); } +SEASTAR_TEST_CASE(to_json_nonprimary_key_eq) { + return do_with_cql_env_thread([](cql_test_env& e) { + cquery_nofail(e, "create table ks.t(pk int, ck int, r int, v vector, primary key(pk, ck))"); + + auto restr = make_restrictions("pk=1 and r=42", e); + auto json = get_restrictions_json(restr, true); + + auto expected = R"json({"restrictions":[{"type":"==","lhs":"pk","rhs":1},{"type":"==","lhs":"r","rhs":42}],"allow_filtering":true})json"; + BOOST_CHECK_EQUAL(json, expected); + }); +} + +SEASTAR_TEST_CASE(to_json_nonprimary_key_range) { + return do_with_cql_env_thread([](cql_test_env& e) { + cquery_nofail(e, "create table ks.t(pk int, ck int, r int, v vector, primary key(pk, ck))"); + + auto restr = make_restrictions("pk=1 and r>10 and r<100", e); + auto json = get_restrictions_json(restr, true); + + auto expected = R"json({"restrictions":[{"type":"==","lhs":"pk","rhs":1},{"type":">","lhs":"r","rhs":10},{"type":"<","lhs":"r","rhs":100}],"allow_filtering":true})json"; + BOOST_CHECK_EQUAL(json, expected); + }); +} + +SEASTAR_TEST_CASE(to_json_nonprimary_key_bind_marker) { + return do_with_cql_env_thread([](cql_test_env& e) { + cquery_nofail(e, "create table ks.t(pk int, ck int, r int, v vector, primary key(pk, ck))"); + + auto restr = make_restrictions("pk=1 and r=?", e); + auto filter = vector_search::prepare_filter(restr, true); + + std::vector bind_values = {raw_value::make_value(int32_type->decompose(99))}; + auto options = make_query_options(std::move(bind_values)); + auto json = rjson::print(filter.to_json(options)); + + auto expected = R"json({"restrictions":[{"type":"==","lhs":"pk","rhs":1},{"type":"==","lhs":"r","rhs":99}],"allow_filtering":true})json"; + BOOST_CHECK_EQUAL(json, expected); + }); +} + BOOST_AUTO_TEST_SUITE_END() diff --git a/vector_search/filter.cc b/vector_search/filter.cc index 9f76101ed5..ee8dadf0a2 100644 --- a/vector_search/filter.cc +++ b/vector_search/filter.cc @@ -186,11 +186,15 @@ prepared_filter prepare_filter(const cql3::restrictions::statement_restrictions& auto& partition_key_restrictions = restrictions.get_partition_key_restrictions(); auto& clustering_columns_restrictions = restrictions.get_clustering_columns_restrictions(); + auto& nonprimary_key_restrictions = restrictions.get_nonprimary_key_restrictions(); expression_to_prepared(partition_key_restrictions, prepared_restrictions); expression_to_prepared(clustering_columns_restrictions, prepared_restrictions); + expression_to_prepared(nonprimary_key_restrictions, prepared_restrictions); - bool has_bind_markers = cql3::expr::contains_bind_marker(partition_key_restrictions) || cql3::expr::contains_bind_marker(clustering_columns_restrictions); + bool has_bind_markers = cql3::expr::contains_bind_marker(partition_key_restrictions) + || cql3::expr::contains_bind_marker(clustering_columns_restrictions) + || cql3::expr::contains_bind_marker(nonprimary_key_restrictions); if (!has_bind_markers) { auto cached_json = restrictions_to_json(prepared_restrictions, allow_filtering, cql3::query_options({})); diff --git a/vector_search/filter.hh b/vector_search/filter.hh index 8912e52a3f..454a0abd39 100644 --- a/vector_search/filter.hh +++ b/vector_search/filter.hh @@ -55,7 +55,7 @@ public: }; /// Prepares a filter from CQL statement restrictions for use in Vector Store service. -/// This function extracts primary key restrictions from the statement_restrictions +/// This function extracts restrictions from the statement_restrictions /// and prepares them for serialization to JSON compatible to Vector Store service filtering API. prepared_filter prepare_filter(const cql3::restrictions::statement_restrictions& restrictions, bool allow_filtering);