From: Yuval Lifshitz Date: Mon, 24 Nov 2025 13:42:52 +0000 (+0000) Subject: rgw/s3vector: add lancedb support X-Git-Url: http://git-server-git.apps.pok.os.sepia.ceph.com/?a=commitdiff_plain;h=d2ed27764620150136fd3feb0d790d0afaa2f217;p=ceph.git rgw/s3vector: add lancedb support not implemented in this commit: * caching information as VectorBucket attributes and fetching them from there: schema, distance type * VectorBucket policy APIs * metadata support Signed-off-by: Yuval Lifshitz --- diff --git a/.gitmodules b/.gitmodules index 58bd31d06bf..ff6505e38b8 100644 --- a/.gitmodules +++ b/.gitmodules @@ -85,3 +85,6 @@ [submodule "src/lss"] path = src/lss url = https://chromium.googlesource.com/linux-syscall-support +[submodule "src/lancedb-c"] + path = src/lancedb-c + url = https://github.com/lancedb/lancedb-c.git diff --git a/CMakeLists.txt b/CMakeLists.txt index 43ce6ef4147..b87dcd563cd 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -597,6 +597,7 @@ option(WITH_RADOSGW_RADOS "RADOS backend for Rados Gateway" ON) option(WITH_RADOSGW_SELECT_PARQUET "Support for s3 select on parquet objects" ON) option(WITH_RADOSGW_ARROW_FLIGHT "Build arrow flight when not using system-provided arrow" OFF) option(WITH_RADOSGW_BACKTRACE_LOGGING "Enable backtraces in rgw logs" OFF) +option(WITH_RADOSGW_LANCEDB "Enable lancedb support in RADOS Gateway" ON) option(WITH_SYSTEM_ARROW "Use system-provided arrow" OFF) option(WITH_SYSTEM_UTF8PROC "Use system-provided utf8proc" OFF) @@ -655,6 +656,13 @@ if(WITH_RADOSGW) endif() message(STATUS "ssl soname: ${LIBSSL_SONAME}") message(STATUS "crypto soname: ${LIBCRYPTO_SONAME}") + if(WITH_RADOSGW_LANCEDB) + # Check for Rust and Cargo + find_program(CARGO_EXECUTABLE cargo) + if(NOT CARGO_EXECUTABLE) + message(FATAL_ERROR "cargo not found. Please install Rust and Cargo to build with WITH_RADOSGW_LANCEDB=ON") + endif() + endif(WITH_RADOSGW_LANCEDB) endif (WITH_RADOSGW) if(WITH_RADOSGW_FDB) diff --git a/ceph.spec.in b/ceph.spec.in index 7bb56a985a8..9d3fdc66d38 100644 --- a/ceph.spec.in +++ b/ceph.spec.in @@ -50,6 +50,7 @@ %bcond_without lttng %bcond_without libradosstriper %bcond_without ocf +%bcond_without lancedb %global luarocks_package_name luarocks %bcond_without lua_packages %global _remote_tarball_prefix https://download.ceph.com/tarballs/ @@ -335,6 +336,9 @@ BuildRequires: librdkafka-devel >= 1.1.0 Requires: lua-devel Requires: %{luarocks_package_name} %endif +%if 0%{with lancedb} +BuildRequires: cargo +%endif %if 0%{with make_check} BuildRequires: hostname BuildRequires: jq @@ -1874,6 +1878,11 @@ cmake .. \ %if 0%{without lua_packages} -DWITH_RADOSGW_LUA_PACKAGES:BOOL=OFF \ %endif +%if 0%{with lancedb} + -DWITH_RADOSGW_LANCEDB:BOOL=ON \ +%else + -DWITH_RADOSGW_LANCEDB:BOOL=OFF \ +%endif %if 0%{with cmake_verbose_logging} -DCMAKE_VERBOSE_MAKEFILE:BOOL=ON \ %endif diff --git a/debian/control b/debian/control index b91fd483239..bc03b68e8fd 100644 --- a/debian/control +++ b/debian/control @@ -9,6 +9,7 @@ Uploaders: Ken Dreyer , Alfredo Deza , Build-Depends: automake, bison, + cargo, cmake (>= 3.10.2), cpio, cython3, diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 6674ccdd38c..23da375960f 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -1200,6 +1200,31 @@ if(WITH_RADOSGW) add_subdirectory(crypto) add_subdirectory(libkmip) + if(WITH_RADOSGW_LANCEDB) + include(ExternalProject) + + set(LANCEDB_BUILD_DIR "${CMAKE_BINARY_DIR}/src/lancedb-c-prefix/src/lancedb-c-build") + # Determine the Rust target directory based on build type + if(CMAKE_BUILD_TYPE STREQUAL "Debug") + set(LANCEDB_RUST_TARGET_DIR "${LANCEDB_BUILD_DIR}/target/debug") + else() + set(LANCEDB_RUST_TARGET_DIR "${LANCEDB_BUILD_DIR}/target/release") + endif() + + ExternalProject_Add(lancedb-c + SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/lancedb-c + CMAKE_ARGS -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE} + -DCMAKE_INSTALL_PREFIX=${CMAKE_BINARY_DIR}/lancedb-c-install + BUILD_ALWAYS OFF + BUILD_BYPRODUCTS ${CMAKE_LIBRARY_OUTPUT_DIRECTORY}/liblancedb.so + INSTALL_COMMAND ${CMAKE_COMMAND} -E copy + ${LANCEDB_RUST_TARGET_DIR}/liblancedb.so + ${CMAKE_LIBRARY_OUTPUT_DIRECTORY}/liblancedb.so + COMMAND ${CMAKE_COMMAND} -E copy + ${CMAKE_CURRENT_SOURCE_DIR}/lancedb-c/include/lancedb.h + ${CMAKE_BINARY_DIR}/src/include/lancedb.h + ) + endif(WITH_RADOSGW_LANCEDB) add_subdirectory(rgw) endif(WITH_RADOSGW) diff --git a/src/lancedb-c b/src/lancedb-c new file mode 160000 index 00000000000..c197134f171 --- /dev/null +++ b/src/lancedb-c @@ -0,0 +1 @@ +Subproject commit c197134f171b6972240ebf7f9d34d874326e824f diff --git a/src/rgw/CMakeLists.txt b/src/rgw/CMakeLists.txt index 2377caf96de..267b9978a03 100644 --- a/src/rgw/CMakeLists.txt +++ b/src/rgw/CMakeLists.txt @@ -411,6 +411,17 @@ if(WITH_RADOSGW_DAOS) link_directories( ${PC_DAOS_LIBRARY_DIRS} ) endif() +if(WITH_RADOSGW_LANCEDB) + # Create imported library target for liblancedb + add_library(lancedb SHARED IMPORTED) + set_target_properties(lancedb PROPERTIES + IMPORTED_LOCATION ${CMAKE_LIBRARY_OUTPUT_DIRECTORY}/liblancedb.so + IMPORTED_NO_SONAME TRUE + ) + target_link_libraries(rgw_common PUBLIC lancedb) + add_dependencies(rgw_common lancedb-c) +endif() + set(rgw_a_srcs rgw_appmain.cc rgw_asio_client.cc diff --git a/src/rgw/rgw_rest_s3vector.cc b/src/rgw/rgw_rest_s3vector.cc index 10aac50112b..53e1d8e7f29 100644 --- a/src/rgw/rgw_rest_s3vector.cc +++ b/src/rgw/rgw_rest_s3vector.cc @@ -6,6 +6,7 @@ #include "rgw_s3vector.h" #include "rgw_process_env.h" #include "common/async/yield_context.h" +#include "rgw_arn.h" #define dout_context g_ceph_context #define dout_subsys ceph_subsys_rgw @@ -49,7 +50,7 @@ class RGWS3VectorCreateIndex : public RGWS3VectorBase { int verify_permission(optional_yield y) override { ldpp_dout(this, 10) << "INFO: verifying permission for s3vector CreateIndex" << dendl; - // TODO: implement permission check + // policy TODO: implement permission check /*if (!verify_bucket_permission(this, s, rgw::IAM::s3vectorsCreateIndex)) { return -EACCES; }*/ @@ -75,6 +76,27 @@ class RGWS3VectorCreateIndex : public RGWS3VectorBase { } op_ret = rgw::s3vector::create_index(configuration, this, y); } + + void send_response() override { + if (op_ret) { + set_req_state_err(s, op_ret); + } + dump_errno(s); + end_header(s, this, "application/json"); + + if (op_ret < 0) { + return; + } + + dump_start(s); + + rgw::s3vector::create_index_reply_t reply{ + rgw::s3vector::index_arn(s->zonegroup_name, s->account_name, configuration.vector_bucket_name, configuration.index_name).to_string() + }; + const auto f = s->formatter; + reply.dump(f); + rgw_flush_formatter_and_reset(s, f); + } }; class RGWS3VectorCreateVectorBucket : public RGWS3VectorBase { @@ -120,7 +142,7 @@ class RGWS3VectorCreateVectorBucket : public RGWS3VectorBase { const rgw_bucket bucket_id(s->bucket_tenant, configuration.vector_bucket_name); int ret = driver->load_vector_bucket(this, bucket_id, &bucket, y); if (ret != -ENOENT) { - // TODO: verify attributes from the existing bucket match the requested configuration + // TODO: verify creation parameters are the same as the existing ones. reject if not op_ret = ret; return; } @@ -179,11 +201,16 @@ class RGWS3VectorCreateVectorBucket : public RGWS3VectorBase { set_req_state_err(s, op_ret); } dump_errno(s); + s->format = RGWFormat::JSON; + end_header(s, this, to_mime_type(s->format)); + if (op_ret < 0) { + return; + } + + dump_start(s); + JSONFormatter f; /* use json formatter for system requests output */ if (op_ret == 0 && s->system_request) { - s->format = RGWFormat::JSON; - end_header(s, this, to_mime_type(s->format)); - JSONFormatter f; /* use json formatter for system requests output */ ceph_assert(bucket); ceph_assert(!bucket->empty()); @@ -194,12 +221,17 @@ class RGWS3VectorCreateVectorBucket : public RGWS3VectorBase { encode_json("object_ver", info.objv_tracker.read_version, &f); encode_json("bucket_info", info, &f); f.close_section(); - rgw_flush_formatter_and_reset(s, &f); - return; } - end_header(s); - } + rgw::s3vector::create_vector_bucket_reply_t reply{ + rgw::s3vector::vector_bucket_arn( + s->zonegroup_name, + s->account_name, + configuration.vector_bucket_name).to_string() + }; + reply.dump(&f); + rgw_flush_formatter_and_reset(s, &f); + } }; class RGWS3VectorDeleteIndex : public RGWS3VectorBase { @@ -207,7 +239,7 @@ class RGWS3VectorDeleteIndex : public RGWS3VectorBase { int verify_permission(optional_yield y) override { ldpp_dout(this, 10) << "INFO: verifying permission for s3vector DeleteIndex" << dendl; - // TODO: implement permission check + // policy TODO: implement permission check /*if (!verify_bucket_permission(this, s, rgw::IAM::s3vectorsDeleteIndex)) { return -EACCES; }*/ @@ -236,15 +268,12 @@ class RGWS3VectorDeleteIndex : public RGWS3VectorBase { }; class RGWS3VectorDeleteVectorBucket : public RGWS3VectorBase { - // TODO: collapse with get_vector_bucket_t an create a common function - // to build and verify the ARN in init_processing() rgw::s3vector::delete_vector_bucket_t configuration; - boost::optional arn; int verify_permission(optional_yield y) override { ldpp_dout(this, 10) << "INFO: verifying permission for s3vector DeleteVectorBucket" << dendl; - if (!verify_bucket_permission(this, s, arn.get(), rgw::IAM::s3vectorsDeleteVectorBucket)) { - // TODO: ignore failure for now "evaluate_iam_policies: implicit deny from identity-based policy" + if (!verify_bucket_permission(this, s, configuration.vector_bucket_arn.get(), rgw::IAM::s3vectorsDeleteVectorBucket)) { + // policy TODO: ignore failure for now "evaluate_iam_policies: implicit deny from identity-based policy" // return -EACCES; return 0; } @@ -261,36 +290,16 @@ class RGWS3VectorDeleteVectorBucket : public RGWS3VectorBase { if (ret < 0) { return ret; } - if (configuration.vector_bucket_arn.empty()) { - arn.emplace(rgw::Partition::aws, - rgw::Service::s3vectors, - s->penv.site->get_zonegroup().api_name, - s->auth.identity->get_tenant(), - configuration.vector_bucket_name); - } else { - arn = rgw::ARN::parse(configuration.vector_bucket_arn); - if (!arn) { - ldpp_dout(this, 1) << "ERROR: invalid s3vector bucket ARN: " << configuration.vector_bucket_arn << dendl; - return -EINVAL; - } - if (arn->service != rgw::Service::s3vectors || - arn->partition != rgw::Partition::aws || - arn->region != s->penv.site->get_zonegroup().api_name || - arn->account != s->bucket_tenant || - arn->resource.empty()) { - ldpp_dout(this, 1) << "ERROR: invalid s3vector bucket ARN service: " << arn->to_string() << dendl; - return -EINVAL; - } - } - if (!configuration.vector_bucket_name.empty() && - arn->resource != configuration.vector_bucket_name) { - ldpp_dout(this, 1) << "ERROR: s3vector bucket ARN bucket mismatch: " << arn->to_string() - << " expected bucket: " << configuration.vector_bucket_name << dendl; - return -EINVAL; + if (!configuration.vector_bucket_arn) { + configuration.vector_bucket_arn = rgw::s3vector::vector_bucket_arn( + s->zonegroup_name, + s->account_name, + configuration.vector_bucket_name + ); } - s->bucket_name = arn->resource; - ldpp_dout(this, 20) << "INFO: s3vector bucket ARN: " << arn.get() << dendl; + s->bucket_name = configuration.vector_bucket_arn->resource; + ldpp_dout(this, 20) << "INFO: s3vector bucket ARN: " << configuration.vector_bucket_arn.get() << dendl; return 0; } @@ -314,8 +323,7 @@ class RGWS3VectorDeleteVectorBucket : public RGWS3VectorBase { ldpp_dout(this, 1) << "ERROR: failed to delete s3vector bucket " << bucket_id << ". error: " << op_ret << dendl; return; } - // TODO: verify bucket is empty before deletion (or support force delete) - // rgw::s3vector::delete_vector_bucket(configuration, this, y); + op_ret = rgw::s3vector::delete_vector_bucket(configuration, this, y); } }; @@ -324,7 +332,7 @@ class RGWS3VectorDeleteVectorBucketPolicy : public RGWS3VectorBase { int verify_permission(optional_yield y) override { ldpp_dout(this, 10) << "INFO: verifying permission for s3vector DeleteVectorBucketPolicy" << dendl; - // TODO: implement permission check + // policy TODO: implement permission check /*if (!verify_bucket_permission(this, s, rgw::IAM::s3vectorsDeleteVectorBucketPolicy)) { return -EACCES; }*/ @@ -357,7 +365,7 @@ class RGWS3VectorPutVectors : public RGWS3VectorBase { int verify_permission(optional_yield y) override { ldpp_dout(this, 10) << "INFO: verifying permission for s3vector PutVectors" << dendl; - // TODO: implement permission check + // policy TODO: implement permission check /*if (!verify_bucket_permission(this, s, rgw::IAM::s3vectorsPutVectors)) { return -EACCES; }*/ @@ -387,10 +395,11 @@ class RGWS3VectorPutVectors : public RGWS3VectorBase { class RGWS3VectorGetVectors : public RGWS3VectorBase { rgw::s3vector::get_vectors_t configuration; + rgw::s3vector::get_vectors_reply_t reply; int verify_permission(optional_yield y) override { ldpp_dout(this, 10) << "INFO: verifying permission for s3vector GetVectors" << dendl; - // TODO: implement permission check + // policy TODO: implement permission check /*if (!verify_bucket_permission(this, s, rgw::IAM::s3vectorsGetVectors)) { return -EACCES; }*/ @@ -414,16 +423,35 @@ class RGWS3VectorGetVectors : public RGWS3VectorBase { ldpp_dout(this, 1) << "ERROR: failed to load s3vector bucket " << bucket_id << ". error: " << op_ret << dendl; return; } - op_ret = rgw::s3vector::get_vectors(configuration, this, y); + op_ret = rgw::s3vector::get_vectors(configuration, this, y, reply); + } + + void send_response() override { + if (op_ret) { + set_req_state_err(s, op_ret); + } + dump_errno(s); + end_header(s, this, "application/json"); + + if (op_ret < 0) { + return; + } + + dump_start(s); + + const auto f = s->formatter; + reply.dump(f); + rgw_flush_formatter_and_reset(s, f); } }; class RGWS3VectorListVectors : public RGWS3VectorBase { rgw::s3vector::list_vectors_t configuration; + rgw::s3vector::list_vectors_reply_t reply; int verify_permission(optional_yield y) override { ldpp_dout(this, 10) << "INFO: verifying permission for s3vector ListVectors" << dendl; - // TODO: implement permission check + // policy TODO: implement permission check /*if (!verify_bucket_permission(this, s, rgw::IAM::s3vectorsListVectors)) { return -EACCES; }*/ @@ -447,7 +475,25 @@ class RGWS3VectorListVectors : public RGWS3VectorBase { ldpp_dout(this, 1) << "ERROR: failed to load s3vector bucket " << bucket_id << ". error: " << op_ret << dendl; return; } - op_ret = rgw::s3vector::list_vectors(configuration, this, y); + op_ret = rgw::s3vector::list_vectors(configuration, this, y, reply); + } + + void send_response() override { + if (op_ret) { + set_req_state_err(s, op_ret); + } + dump_errno(s); + end_header(s, this, "application/json"); + + if (op_ret < 0) { + return; + } + + dump_start(s); + + const auto f = s->formatter; + reply.dump(f); + rgw_flush_formatter_and_reset(s, f); } }; @@ -536,13 +582,12 @@ class RGWS3VectorListVectorBuckets : public RGWS3VectorBase { class RGWS3VectorGetVectorBucket : public RGWS3VectorBase { rgw::s3vector::get_vector_bucket_t configuration; std::unique_ptr bucket; - boost::optional arn; int verify_permission(optional_yield y) override { ldpp_dout(this, 10) << "INFO: verifying permission for s3vector GetVectorBucket" << dendl; - if (!verify_bucket_permission(this, s, arn.get(), rgw::IAM::s3vectorsGetVectorBucket)) { - // TODO: ignore failure for now "evaluate_iam_policies: implicit deny from identity-based policy" + if (!verify_bucket_permission(this, s, configuration.vector_bucket_arn.get(), rgw::IAM::s3vectorsGetVectorBucket)) { + // policy TODO: ignore failure for now "evaluate_iam_policies: implicit deny from identity-based policy" // return -EACCES; return 0; } @@ -559,36 +604,16 @@ class RGWS3VectorGetVectorBucket : public RGWS3VectorBase { if (ret < 0) { return ret; } - if (configuration.vector_bucket_arn.empty()) { - arn.emplace(rgw::Partition::aws, - rgw::Service::s3vectors, - s->penv.site->get_zonegroup().api_name, - s->auth.identity->get_tenant(), - configuration.vector_bucket_name); - } else { - arn = rgw::ARN::parse(configuration.vector_bucket_arn); - if (!arn) { - ldpp_dout(this, 1) << "ERROR: invalid s3vector bucket ARN: " << configuration.vector_bucket_arn << dendl; - return -EINVAL; - } - if (arn->service != rgw::Service::s3vectors || - arn->partition != rgw::Partition::aws || - arn->region != s->penv.site->get_zonegroup().api_name || - arn->account != s->bucket_tenant || - arn->resource.empty()) { - ldpp_dout(this, 1) << "ERROR: invalid s3vector bucket ARN service: " << arn->to_string() << dendl; - return -EINVAL; - } - } - if (!configuration.vector_bucket_name.empty() && - arn->resource != configuration.vector_bucket_name) { - ldpp_dout(this, 1) << "ERROR: s3vector bucket ARN bucket mismatch: " << arn->to_string() - << " expected bucket: " << configuration.vector_bucket_name << dendl; - return -EINVAL; + if (!configuration.vector_bucket_arn) { + configuration.vector_bucket_arn = rgw::s3vector::vector_bucket_arn( + s->zonegroup_name, + s->account_name, + configuration.vector_bucket_name + ); } - s->bucket_name = arn->resource; - ldpp_dout(this, 20) << "INFO: s3vector bucket ARN: " << arn.get() << dendl; + s->bucket_name = configuration.vector_bucket_arn->resource; + ldpp_dout(this, 20) << "INFO: s3vector bucket ARN: " << configuration.vector_bucket_arn.get() << dendl; return 0; } @@ -618,9 +643,7 @@ class RGWS3VectorGetVectorBucket : public RGWS3VectorBase { f.open_object_section(""); f.open_object_section("vectorBucket"); ::encode_json("creationTime", ceph::to_iso_8601(bucket->get_creation_time()), &f); - rgw::ARN arn(rgw::Partition::aws, rgw::Service::s3vectors, - "", bucket->get_tenant(), bucket->get_name()); - ::encode_json("vectorBucketArn", arn.to_string(), &f); + ::encode_json("vectorBucketArn", configuration.vector_bucket_arn->to_string(), &f); ::encode_json("vectorBucketName", bucket->get_name(), &f); f.close_section(); // vectorBucket f.close_section(); // root object @@ -632,10 +655,11 @@ class RGWS3VectorGetVectorBucket : public RGWS3VectorBase { class RGWS3VectorGetIndex : public RGWS3VectorBase { rgw::s3vector::get_index_t configuration; + rgw::s3vector::get_index_reply_t reply; int verify_permission(optional_yield y) override { ldpp_dout(this, 10) << "INFO: verifying permission for s3vector GetIndex" << dendl; - // TODO: implement permission check + // policy TODO: implement permission check /*if (!verify_bucket_permission(this, s, rgw::IAM::s3vectorsGetIndex)) { return -EACCES; }*/ @@ -659,16 +683,35 @@ class RGWS3VectorGetIndex : public RGWS3VectorBase { ldpp_dout(this, 1) << "ERROR: failed to load s3vector bucket " << bucket_id << ". error: " << op_ret << dendl; return; } - op_ret = rgw::s3vector::get_index(configuration, this, y); + op_ret = rgw::s3vector::get_index(configuration, s->zonegroup_name, s->account_name, this, y, reply); + } + + void send_response() override { + if (op_ret) { + set_req_state_err(s, op_ret); + } + dump_errno(s); + end_header(s, this, "application/json"); + + if (op_ret < 0) { + return; + } + + dump_start(s); + + const auto f = s->formatter; + reply.dump(f); + rgw_flush_formatter_and_reset(s, f); } }; class RGWS3VectorListIndexes : public RGWS3VectorBase { rgw::s3vector::list_indexes_t configuration; + rgw::s3vector::list_indexes_reply_t reply; int verify_permission(optional_yield y) override { ldpp_dout(this, 10) << "INFO: verifying permission for s3vector ListIndexes" << dendl; - // TODO: implement permission check + // policy TODO: implement permission check /*if (!verify_bucket_permission(this, s, rgw::IAM::s3vectorsListIndexes)) { return -EACCES; }*/ @@ -692,7 +735,32 @@ class RGWS3VectorListIndexes : public RGWS3VectorBase { ldpp_dout(this, 1) << "ERROR: failed to load s3vector bucket " << bucket_id << ". error: " << op_ret << dendl; return; } - op_ret = rgw::s3vector::list_indexes(configuration, this, y); + if (!configuration.vector_bucket_arn) { + configuration.vector_bucket_arn = rgw::s3vector::vector_bucket_arn( + s->zonegroup_name, + s->account_name, + configuration.vector_bucket_name + ); + } + op_ret = rgw::s3vector::list_indexes(configuration, this, y, reply); + } + + void send_response() override { + if (op_ret) { + set_req_state_err(s, op_ret); + } + dump_errno(s); + end_header(s, this, "application/json"); + + if (op_ret < 0) { + return; + } + + dump_start(s); + + const auto f = s->formatter; + reply.dump(f); + rgw_flush_formatter_and_reset(s, f); } }; @@ -701,7 +769,7 @@ class RGWS3VectorPutVectorBucketPolicy : public RGWS3VectorBase { int verify_permission(optional_yield y) override { ldpp_dout(this, 10) << "INFO: verifying permission for s3vector PutVectorBucketPolicy" << dendl; - // TODO: implement permission check + // policy TODO: implement permission check /*if (!verify_bucket_permission(this, s, rgw::IAM::s3vectorsPutVectorBucketPolicy)) { return -EACCES; }*/ @@ -725,7 +793,7 @@ class RGWS3VectorPutVectorBucketPolicy : public RGWS3VectorBase { ldpp_dout(this, 1) << "ERROR: failed to load s3vector bucket " << bucket_id << ". error: " << op_ret << dendl; return; } - op_ret = rgw::s3vector::put_vector_bucket_policy(configuration, this, y); + // policy TODO: implement } }; @@ -734,7 +802,7 @@ class RGWS3VectorGetVectorBucketPolicy : public RGWS3VectorBase { int verify_permission(optional_yield y) override { ldpp_dout(this, 10) << "INFO: verifying permission for s3vector GetVectorBucketPolicy" << dendl; - // TODO: implement permission check + // policy TODO: implement permission check /*if (!verify_bucket_permission(this, s, rgw::IAM::s3vectorsGetVectorBucketPolicy)) { return -EACCES; }*/ @@ -758,7 +826,7 @@ class RGWS3VectorGetVectorBucketPolicy : public RGWS3VectorBase { ldpp_dout(this, 1) << "ERROR: failed to load s3vector bucket " << bucket_id << ". error: " << op_ret << dendl; return; } - op_ret = rgw::s3vector::get_vector_bucket_policy(configuration, this, y); + // policy TODO: implement } }; @@ -767,7 +835,7 @@ class RGWS3VectorDeleteVectors : public RGWS3VectorBase { int verify_permission(optional_yield y) override { ldpp_dout(this, 10) << "INFO: verifying permission for s3vector DeleteVectors" << dendl; - // TODO: implement permission check + // policy TODO: implement permission check /*if (!verify_bucket_permission(this, s, rgw::IAM::s3vectorsDeleteVectors)) { return -EACCES; }*/ @@ -797,10 +865,11 @@ class RGWS3VectorDeleteVectors : public RGWS3VectorBase { class RGWS3VectorQueryVectors : public RGWS3VectorBase { rgw::s3vector::query_vectors_t configuration; + rgw::s3vector::query_vectors_reply_t reply; int verify_permission(optional_yield y) override { ldpp_dout(this, 10) << "INFO: verifying permission for s3vector QueryVectors" << dendl; - // TODO: implement permission check + // policy TODO: implement permission check /*if (!verify_bucket_permission(this, s, rgw::IAM::s3vectorsQueryVectors)) { return -EACCES; }*/ @@ -824,9 +893,36 @@ class RGWS3VectorQueryVectors : public RGWS3VectorBase { ldpp_dout(this, 1) << "ERROR: failed to load s3vector bucket " << bucket_id << ". error: " << op_ret << dendl; return; } - op_ret = rgw::s3vector::query_vectors(configuration, this, y); + op_ret = rgw::s3vector::query_vectors(configuration, this, y, reply); + } + + void send_response() override { + if (op_ret) { + set_req_state_err(s, op_ret); + } + dump_errno(s); + end_header(s, this, "application/json"); + + if (op_ret < 0) { + return; + } + + dump_start(s); + + const auto f = s->formatter; + reply.dump(f); + rgw_flush_formatter_and_reset(s, f); } }; + +} + +int RGWHandler_REST_s3Vector::init(rgw::sal::Driver* driver, req_state *s, rgw::io::BasicClient *cio) { + if (int ret = RGWHandler_REST_S3::init(driver, s, cio); ret < 0) { + return ret; + } + // s3vectors API uses JSON formatter + return RGWHandler_REST::allocate_formatter(s, RGWFormat::JSON, false); } RGWOp* RGWHandler_REST_s3Vector::op_post() { diff --git a/src/rgw/rgw_rest_s3vector.h b/src/rgw/rgw_rest_s3vector.h index 4b1b3c92c6b..874fde1642c 100644 --- a/src/rgw/rgw_rest_s3vector.h +++ b/src/rgw/rgw_rest_s3vector.h @@ -10,10 +10,13 @@ protected: int init_permissions(RGWOp* op, optional_yield y) override {return 0;} int read_permissions(RGWOp* op, optional_yield y) override {return 0;} bool supports_quota() override {return false;} + public: explicit RGWHandler_REST_s3Vector(const rgw::auth::StrategyRegistry& auth_registry) : RGWHandler_REST_S3(auth_registry) {} virtual ~RGWHandler_REST_s3Vector() = default; + + int init(rgw::sal::Driver* driver, req_state *s, rgw::io::BasicClient *cio) override; RGWOp *op_post() override; static RGWOp* create_post_op(const std::string& op_name); }; diff --git a/src/rgw/rgw_s3vector.cc b/src/rgw/rgw_s3vector.cc index e1fe0754b18..fe449ada98b 100644 --- a/src/rgw/rgw_s3vector.cc +++ b/src/rgw/rgw_s3vector.cc @@ -5,590 +5,1605 @@ #include "common/ceph_json.h" #include "common/Formatter.h" #include "common/dout.h" +#include #include -#include "rgw_sal.h" +#include "lancedb.h" +#include +#include +#include #define dout_subsys ceph_subsys_rgw namespace rgw::s3vector { + // convert LanceDBError to linux error codes + int lancedb_error_to_errno(LanceDBError err) { + switch (err) { + case LANCEDB_SUCCESS: + return 0; + case LANCEDB_INVALID_ARGUMENT: + case LANCEDB_INVALID_TABLE_NAME: + case LANCEDB_INVALID_INPUT: + return -EINVAL; + case LANCEDB_TABLE_NOT_FOUND: + case LANCEDB_DATABASE_NOT_FOUND: + case LANCEDB_INDEX_NOT_FOUND: + case LANCEDB_EMBEDDING_FUNCTION_NOT_FOUND: + return -ENOENT; + case LANCEDB_DATABASE_ALREADY_EXISTS: + case LANCEDB_TABLE_ALREADY_EXISTS: + return -EEXIST; + case LANCEDB_NOT_SUPPORTED: + return -EOPNOTSUPP; + case LANCEDB_RETRY: + return -EAGAIN; + case LANCEDB_TIMEOUT: + return -EBUSY; + case LANCEDB_CREATE_DIR: + case LANCEDB_SCHEMA: + case LANCEDB_RUNTIME: + case LANCEDB_OBJECT_STORE: + case LANCEDB_LANCE: + case LANCEDB_HTTP: + case LANCEDB_ARROW: + case LANCEDB_OTHER: + case LANCEDB_UNKNOWN: + return -EIO; + } + return -EIO; + } -// utility functions for JSON encoding/decoding + // utility functions for connection creation and opening table -void decode_json_obj(float& val, JSONObj *obj) { - std::string s = obj->get_data(); - const char *start = s.c_str(); - char *p; + LanceDBConnection* connect(DoutPrefixProvider* dpp, const std::string& vector_bucket_name) { + const auto dbname = fmt::format("/tmp/lancedb/{}", vector_bucket_name); + LanceDBConnectBuilder* builder = lancedb_connect(dbname.c_str()); + LanceDBConnection* conn = lancedb_connect_builder_execute(builder); + if (!conn) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to connect to: " << dbname << dendl; + } + return conn; + } - errno = 0; - val = strtof(start, &p); + LanceDBTable* open_table(DoutPrefixProvider* dpp, const std::string& vector_bucket_name, const std::string& index_name) { + LanceDBConnection* conn = connect(dpp, vector_bucket_name); + if (!conn) { + return nullptr; + } + LanceDBTable* table = lancedb_connection_open_table(conn, index_name.c_str()); + if (!table) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to open index: " << index_name << " in: " << vector_bucket_name << dendl; + lancedb_connection_free(conn); + return nullptr; + } + return table; + } - /* Check for various possible errors */ + // get creation time from the first version of a table + // returns 0 on failure + uint64_t get_table_creation_time(const LanceDBTable* table, DoutPrefixProvider* dpp) { + LanceDBVersion* versions = nullptr; + size_t count = 0; + char* error_message = nullptr; + if (const auto result = lancedb_table_list_versions(table, &versions, nullptr, &count, &error_message); result != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to list table versions: " << (error_message ? error_message : "unknown") << dendl; + lancedb_free_string(error_message); + return 0; + } + if (count == 0) { + lancedb_free_versions(versions, count); + return 0; + } + const uint64_t creation_time = versions[0].timestamp_seconds; + lancedb_free_versions(versions, count); + return creation_time; + } - if ((errno == ERANGE && val == HUGE_VAL) || - (errno != 0 && val == 0)) { - throw JSONDecoder::err("failed to number"); - } + // utility functions for JSON encoding/decoding - if (p == start) { - throw JSONDecoder::err("failed to parse number"); - } + void decode_json_obj(float& val, JSONObj *obj) { + std::string_view s = obj->get_data(); + const char* start = s.data(); + const char* end = start + s.length(); - while (*p != '\0') { - if (!isspace(*p)) { - throw JSONDecoder::err("failed to parse number"); - } - p++; - } -} + const auto result = std::from_chars(start, end, val); -void decode_json(const char* field_name, DistanceMetric& metric, JSONObj* obj, bool mandatory) { - std::string metric_str; - JSONDecoder::decode_json(field_name, metric_str, obj, mandatory); - if (metric_str == "cosine") { - metric = DistanceMetric::COSINE; - } else if (metric_str == "euclidean") { - metric = DistanceMetric::EUCLIDEAN; - } else { - throw JSONDecoder::err("invalid distanceMetric: " + metric_str); + if (result.ec == std::errc::invalid_argument) { + throw JSONDecoder::err("failed to parse number"); + } + + if (result.ec == std::errc::result_out_of_range) { + throw JSONDecoder::err("out of range number"); + } } -} -void encode_json(const char* field_name, const DistanceMetric& metric, ceph::Formatter* f) { - switch (metric) { - case DistanceMetric::COSINE: - ::encode_json(field_name, "cosine", f); - return; - case DistanceMetric::EUCLIDEAN: - ::encode_json(field_name, "euclidean", f); - return; + void decode_json(const char* field_name, VectorData& data, JSONObj* obj) { + data.clear(); + auto it = obj->find(field_name); + if (it.end()) { + throw JSONDecoder::err(std::string("missing field: ") + field_name); + } + auto arr_it = (*it)->find("float32"); + for (auto value_it = (*arr_it)->find_first(); !value_it.end(); ++value_it) { + float value; + decode_json_obj(value, *value_it); + data.push_back(value); + } } - ::encode_json(field_name, "unknown", f); -} -void decode_json(const char* field_name, VectorData& data, JSONObj* obj) { - data.clear(); - auto it = obj->find(field_name); - if (it.end()) { - throw JSONDecoder::err(std::string("missing field: ") + field_name); + void encode_json(const char* field_name, const VectorData& data, ceph::Formatter* f) { + f->open_object_section(field_name); + f->open_array_section("float32"); + for (auto value : data) { + f->dump_float("", value); + } + f->close_section(); + f->close_section(); } - auto arr_it = (*it)->find("float32"); - for (auto value_it = (*arr_it)->find_first(); !value_it.end(); ++value_it) { - float value; - decode_json_obj(value, *value_it); - data.push_back(value); + + void verify_name(const std::string& name) { + if (name.length() < 3 || name.length() > 63) { + throw JSONDecoder::err(fmt::format("'{}' length ({}) must be between 3 and 63 characters", name, name.length())); + } + } + + void decode_name_or_arn(const char* name_field, const char* arn_field, std::string& name, boost::optional& arn, JSONObj* obj) { + std::string str_arn; + JSONDecoder::decode_json(arn_field, str_arn, obj); + JSONDecoder::decode_json(name_field, name, obj); + if (str_arn.empty() && name.empty()) { + throw JSONDecoder::err(fmt::format("either {} or {} must be specified", name_field, arn_field)); + } + if (!str_arn.empty() && !name.empty()) { + throw JSONDecoder::err(fmt::format("only one of {} and {} must be specified", name_field, arn_field)); + } + if (!str_arn.empty()) { + arn = rgw::ARN::parse(str_arn); + if (!arn) throw JSONDecoder::err(fmt::format("failed to parse ARN {}", str_arn)); + } } -} -void encode_json(const char* field_name, const VectorData& data, ceph::Formatter* f) { - f->open_object_section(field_name); - f->open_array_section("float32"); - for (auto value : data) { - f->dump_float("", value); + void decode_vector_bucket_name(std::string& vector_bucket_name, boost::optional& arn, JSONObj* obj) { + decode_name_or_arn("vectorBucketName", "vectorBucketArn", vector_bucket_name, arn, obj); + if (arn) { + ceph_assert(vector_bucket_name.empty()); + constexpr std::string_view prefix = "bucket/"; + std::string_view resource{arn->resource}; + if (!resource.starts_with(prefix)) { + throw JSONDecoder::err( + fmt::format("invalid vector bucket ARN. expected: 'bucket/' got: {}", resource)); + } + vector_bucket_name = resource.substr(prefix.size()); + } + verify_name(vector_bucket_name); } - f->close_section(); - f->close_section(); -} -void create_index_t::dump(ceph::Formatter* f) const { - f->open_object_section(""); - ::encode_json("dataType", data_type, f); - ::encode_json("dimension", dimension, f); - rgw::s3vector::encode_json("distanceMetric", distance_metric, f); - ::encode_json("indexName", index_name, f); - f->open_object_section("metadataConfiguration"); - ::encode_json("nonFilterableMetadataKeys", non_filterable_metadata_keys, f); - f->close_section(); - ::encode_json("vectorBucketArn", vector_bucket_arn, f); - ::encode_json("vectorBucketName", vector_bucket_name, f); - f->close_section(); -} + void decode_index_name(std::string& vector_bucket_name, std::string& index_name, JSONObj* obj) { + boost::optional arn; + decode_name_or_arn("indexName", "indexArn", index_name, arn, obj); + if (arn) { + ceph_assert(index_name.empty()); + constexpr std::string_view bucket_prefix = "bucket/"; + constexpr std::string_view index_prefix = "/index/"; + std::string_view resource{arn->resource}; + if (!resource.starts_with(bucket_prefix)) { + throw JSONDecoder::err( + fmt::format("invalid index ARN. expected: 'bucket//index/' got: {}", resource)); + } + // Remove "bucket/" prefix + resource.remove_prefix(bucket_prefix.size()); + // Find "/index/" separator + auto index_pos = resource.find(index_prefix); + if (index_pos == std::string_view::npos) { + throw JSONDecoder::err( + fmt::format("invalid index ARN. expected: 'bucket//index/' got: {}", resource)); + } + vector_bucket_name = resource.substr(0, index_pos); + index_name = resource.substr(index_pos + index_prefix.size()); + } else { + JSONDecoder::decode_json("vectorBucketName", vector_bucket_name, obj, true); + } + verify_name(index_name); + verify_name(vector_bucket_name); + } -void decode_name(const char* name_field, std::string& name, JSONObj* obj) { - JSONDecoder::decode_json(name_field, name, obj, true); - if (name.length() < 3 || name.length() > 63) { - throw JSONDecoder::err(fmt::format("{} length must be between 3 and 63 characters, got {}", name_field, name.length())); + template + void log_configuration(DoutPrefixProvider* dpp, const std::string& op_name, const T& configuration) { + JSONFormatter f; + configuration.dump(&f); + std::stringstream ss; + f.flush(ss); + ldpp_dout(dpp, 20) << "INFO: executing s3vector " << op_name << " with: " << ss.str() << dendl; } -} -void decode_name_or_arn(const char* name_field, const char* arn_field, std::string& name, std::string& arn, JSONObj* obj) { - JSONDecoder::decode_json(arn_field, arn, obj); - JSONDecoder::decode_json(name_field, name, obj); - if (arn.empty() && name.empty()) { - throw JSONDecoder::err(fmt::format("either {} or {} must be specified", name_field, arn_field)); + // index operations: create, delete, get, list + ////////////////////////////////////////////// + + // create index + + const char* distance_metric_key = "distance_metric"; + + const char* distance_metric_to_string(DistanceMetric metric) { + switch (metric) { + case DistanceMetric::COSINE: return "cosine"; + case DistanceMetric::EUCLIDEAN: return "euclidean"; + case DistanceMetric::UNKNOWN: return "unknown"; + } + return "unknown"; } - if (!name.empty() && (name.length() < 3 || name.length() > 63)) { - throw JSONDecoder::err(fmt::format("{} length must be between 3 and 63 characters, got {}", name_field, name.length())); + + DistanceMetric string_to_distance_metric(const std::string& str) { + if (str == "cosine") { + return DistanceMetric::COSINE; + } else if (str == "euclidean") { + return DistanceMetric::EUCLIDEAN; + } + return DistanceMetric::UNKNOWN; } - //TODO: validate ARN -} -void create_index_t::decode_json(JSONObj* obj) { - JSONDecoder::decode_json("dataType", data_type, obj, true); - if (data_type != "float32") { - throw JSONDecoder::err(fmt::format("invalid dataType: {}. Only 'float32' is supported.", data_type)); + int set_table_distance_metric(const LanceDBTable* table, DistanceMetric metric, DoutPrefixProvider* dpp) { + const char* key = distance_metric_key; + const char* value = distance_metric_to_string(metric); + char* error_message = nullptr; + if (const auto result = lancedb_table_set_metadata(table, &key, &value, 1, &error_message); result != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to set distance_metric metadata: " << (error_message ? error_message : "unknown") << dendl; + lancedb_free_string(error_message); + return lancedb_error_to_errno(result); + } + return 0; } - JSONDecoder::decode_json("dimension", dimension, obj, true); - if (dimension < 1 || dimension > 4096) { - throw JSONDecoder::err(fmt::format("dimension must be between 1 and 4096, got {}", dimension)); + + DistanceMetric get_table_distance_metric(const LanceDBTable* table, DoutPrefixProvider* dpp) { + const char* key = distance_metric_key; + char** keys_out = nullptr; + char** values_out = nullptr; + size_t count = 0; + char* error_message = nullptr; + if (const auto result = lancedb_table_get_metadata(table, &key, 1, &keys_out, &values_out, &count, &error_message); result != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to get distance_metric metadata: " << (error_message ? error_message : "unknown") << dendl; + lancedb_free_string(error_message); + return DistanceMetric::UNKNOWN; + } + DistanceMetric metric = DistanceMetric::UNKNOWN; + if (count > 0) { + metric = string_to_distance_metric(values_out[0]); + } + lancedb_free_metadata(keys_out, values_out, count); + return metric; } - rgw::s3vector::decode_json("distanceMetric", distance_metric, obj, true); - JSONDecoder::decode_json("indexName", index_name, obj, true); - auto md_it = obj->find("metadataConfiguration"); - if (!md_it.end()) { - JSONDecoder::decode_json("nonFilterableMetadataKeys", non_filterable_metadata_keys, *md_it); + + void create_index_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + ::encode_json("dataType", data_type, f); + ::encode_json("dimension", dimension, f); + ::encode_json("distanceMetric", distance_metric_to_string(distance_metric), f); + ::encode_json("indexName", index_name, f); + f->open_object_section("metadataConfiguration"); + ::encode_json("nonFilterableMetadataKeys", non_filterable_metadata_keys, f); + f->close_section(); + if (vector_bucket_arn) { + ::encode_json("vectorBucketArn", vector_bucket_arn->to_string(), f); + } else { + ::encode_json("vectorBucketName", vector_bucket_name, f); + } + f->close_section(); } - decode_name_or_arn("vectorBucketName", "vectorBucketArn", vector_bucket_name, vector_bucket_arn, obj); -} -void create_vector_bucket_t::dump(ceph::Formatter* f) const { - f->open_object_section(""); - ::encode_json("vectorBucketName", vector_bucket_name, f); - f->close_section(); -} + void create_index_t::decode_json(JSONObj* obj) { + JSONDecoder::decode_json("dataType", data_type, obj, true); + if (data_type != "float32") { + throw JSONDecoder::err(fmt::format("invalid dataType: {}. Only 'float32' is supported.", data_type)); + } + JSONDecoder::decode_json("dimension", dimension, obj, true); + if (dimension < 1 || dimension > 4096) { + throw JSONDecoder::err(fmt::format("dimension must be between 1 and 4096, got {}", dimension)); + } + std::string metric_str; + JSONDecoder::decode_json("distanceMetric", metric_str, obj, true); + distance_metric = string_to_distance_metric(metric_str); + if (distance_metric == DistanceMetric::UNKNOWN) { + throw JSONDecoder::err("invalid distanceMetric: " + metric_str); + } + JSONDecoder::decode_json("indexName", index_name, obj, true); + auto md_it = obj->find("metadataConfiguration"); + if (!md_it.end()) { + JSONDecoder::decode_json("nonFilterableMetadataKeys", non_filterable_metadata_keys, *md_it); + } + decode_vector_bucket_name(vector_bucket_name, vector_bucket_arn, obj); + } -void create_vector_bucket_t::decode_json(JSONObj* obj) { - decode_name("vectorBucketName", vector_bucket_name, obj); -} + void create_index_reply_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + ::encode_json("indexArn", index_arn, f); + f->close_section(); + } -int create_index(const create_index_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { - JSONFormatter f; - configuration.dump(&f); - std::stringstream ss; - f.flush(ss); - ldpp_dout(dpp, 20) << "INFO: executing s3vector CreateIndex with: " << ss.str() << dendl; + static constexpr const char* data_field = "data"; + static const std::string data_field_str{data_field};; + static constexpr const char* key_field = "key"; + static const std::string key_field_str{key_field};; + static constexpr const char* distance_field = "_distance"; + static const std::string distance_field_str{distance_field};; + static constexpr const char* key_columns[] = {key_field}; + static constexpr const char* data_columns[] = {data_field}; + static constexpr const char* table_columns[] = {key_field, data_field}; + static constexpr int num_key_columns = 1; + static constexpr int num_table_columns = 2; + + int create_table_schema(unsigned int dimension, DoutPrefixProvider* dpp, ArrowSchema* c_schema) { + const auto schema = arrow::schema( + { + arrow::field(key_field, arrow::utf8()), + arrow::field(data_field, arrow::fixed_size_list(arrow::float32(), dimension)) + } + ); + if (const auto status = arrow::ExportSchema(*schema, c_schema); !status.ok()) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to export schema to C ABI: " << status.ToString() << dendl; + return -EINVAL; + } return 0; -} + } -void delete_index_t::dump(ceph::Formatter* f) const { - f->open_object_section(""); - ::encode_json("indexArn", index_arn, f); - ::encode_json("indexName", index_name, f); - ::encode_json("vectorBucketName", vector_bucket_name, f); - f->close_section(); -} + int get_vector_dimension(const std::string& index_name, LanceDBTable* table, DoutPrefixProvider* dpp, unsigned int& dimension); -void delete_index_t::decode_json(JSONObj* obj) { - decode_name_or_arn("indexName", "indexArn", index_name, index_arn, obj); - decode_name("vectorBucketName", vector_bucket_name, obj); -} + int create_index(const create_index_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { + log_configuration(dpp, "CreateIndex", configuration); + LanceDBConnection* conn = connect(dpp, configuration.vector_bucket_name); + if (!conn) { + return -EIO; + } + // if the table already exists, verify the schema matches + LanceDBTable* existing_table = lancedb_connection_open_table(conn, configuration.index_name.c_str()); + if (existing_table) { + unsigned int existing_dimension = 0; + if (get_vector_dimension(configuration.index_name, existing_table, dpp, existing_dimension) < 0) { + lancedb_table_free(existing_table); + lancedb_connection_free(conn); + return -EEXIST; + } + if (existing_dimension != configuration.dimension) { + ldpp_dout(dpp, 1) << "ERROR: s3vector index: " << configuration.index_name + << " already exists with dimension " << existing_dimension + << " but requested dimension " << configuration.dimension << dendl; + lancedb_table_free(existing_table); + lancedb_connection_free(conn); + return -EEXIST; + } + const auto existing_metric = get_table_distance_metric(existing_table, dpp); + if (existing_metric != DistanceMetric::UNKNOWN && existing_metric != configuration.distance_metric) { + ldpp_dout(dpp, 1) << "ERROR: s3vector index: " << configuration.index_name + << " already exists with distance metric " << distance_metric_to_string(existing_metric) + << " but requested " << distance_metric_to_string(configuration.distance_metric) << dendl; + lancedb_table_free(existing_table); + lancedb_connection_free(conn); + return -EEXIST; + } + ldpp_dout(dpp, 10) << "INFO: s3vector index: " << configuration.index_name + << " already exists with matching schema, returning success" << dendl; + lancedb_table_free(existing_table); + lancedb_connection_free(conn); + return 0; + } -void delete_vector_bucket_t::dump(ceph::Formatter* f) const { - f->open_object_section(""); - ::encode_json("vectorBucketArn", vector_bucket_arn, f); - ::encode_json("vectorBucketName", vector_bucket_name, f); - f->close_section(); -} + struct ArrowSchema c_schema; + if (int ret = create_table_schema(configuration.dimension, dpp, &c_schema); ret < 0) { + lancedb_connection_free(conn); + return ret; + } + char* error_message; + LanceDBTable* table = nullptr; + if (const LanceDBError result = lancedb_table_create(conn, configuration.index_name.c_str(), + reinterpret_cast(&c_schema), + nullptr, &table, &error_message); result != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: s3vector creating index: " << configuration.index_name << ", error: " << error_message << dendl; + lancedb_free_string(error_message); + lancedb_connection_free(conn); + return lancedb_error_to_errno(result); + } + // create the main index on the table (vector index will be created only after vectors are added) + const LanceDBScalarIndexConfig scalar_config = { + .replace = 1, // replace existing index + .force_update_statistics = 0 // don't force update statistics + }; + if (const LanceDBError result = lancedb_table_create_scalar_index( + table, key_columns, num_key_columns, LANCEDB_INDEX_BTREE, &scalar_config, &error_message); result != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: s3vector creating scalar index: " << configuration.index_name << "on 'key' columns, error: " << error_message << dendl; + lancedb_table_free(table); + lancedb_free_string(error_message); + lancedb_connection_free(conn); + return lancedb_error_to_errno(result); + } + if (int ret = set_table_distance_metric(table, configuration.distance_metric, dpp); ret < 0) { + lancedb_table_free(table); + lancedb_connection_free(conn); + return ret; + } + lancedb_table_free(table); + lancedb_connection_free(conn); + return 0; + } -void delete_vector_bucket_t::decode_json(JSONObj* obj) { - decode_name_or_arn("vectorBucketName", "vectorBucketArn", vector_bucket_name, vector_bucket_arn, obj); -} + // delete index -void delete_vector_bucket_policy_t::dump(ceph::Formatter* f) const { - f->open_object_section(""); - ::encode_json("vectorBucketName", vector_bucket_name, f); - f->close_section(); -} + void delete_index_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + if (index_arn) { + ::encode_json("indexArn", index_arn->to_string(), f); + } else { + ::encode_json("indexName", index_name, f); + ::encode_json("vectorBucketName", vector_bucket_name, f); + } + f->close_section(); + } -void delete_vector_bucket_policy_t::decode_json(JSONObj* obj) { - decode_name_or_arn("vectorBucketName", "vectorBucketArn", vector_bucket_name, vector_bucket_arn, obj); -} + void delete_index_t::decode_json(JSONObj* obj) { + decode_index_name(vector_bucket_name, index_name, obj); + } -int create_vector_bucket(const create_vector_bucket_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { + int delete_index(const delete_index_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { + log_configuration(dpp, "DeleteIndex", configuration); + LanceDBConnection* conn = connect(dpp, configuration.vector_bucket_name); + if (!conn) { + return -EIO; + } + char* error_message; + if (const LanceDBError result = lancedb_connection_drop_table(conn, + configuration.index_name.c_str(), nullptr, &error_message); result != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to delete index: " << configuration.index_name << ". error: " << error_message << dendl; + lancedb_free_string(error_message); + lancedb_connection_free(conn); + return lancedb_error_to_errno(result); + } return 0; -} + } -int delete_index(const delete_index_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { - JSONFormatter f; - configuration.dump(&f); - std::stringstream ss; - f.flush(ss); - ldpp_dout(dpp, 20) << "INFO: executing s3vector DeleteIndex with: " << ss.str() << dendl; - return 0; -} + // get index + + void get_index_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + if (index_arn) { + ::encode_json("indexArn", index_arn->to_string(), f); + } else { + ::encode_json("indexName", index_name, f); + ::encode_json("vectorBucketName", vector_bucket_name, f); + } + f->close_section(); + } + + void get_index_t::decode_json(JSONObj* obj) { + decode_index_name(vector_bucket_name, index_name, obj); + } + + void get_index_reply_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + f->open_object_section("index"); + ::encode_json("creationTime", creation_time, f); + ::encode_json("dataType", data_type, f); + ::encode_json("dimension", dimension, f); + ::encode_json("distanceMetric", distance_metric_to_string(distance_metric), f); + ::encode_json("indexArn", index_arn, f); + ::encode_json("indexName", index_name, f); + f->open_object_section("metadataConfiguration"); + ::encode_json("nonFilterableMetadataKeys", non_filterable_metadata_keys, f); + f->close_section(); + ::encode_json("vectorBucketName", vector_bucket_name, f); + f->close_section(); + f->close_section(); + } -int delete_vector_bucket(const delete_vector_bucket_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { + int get_vector_dimension(const std::string& index_name, LanceDBTable* table, DoutPrefixProvider* dpp, unsigned int& dimension) { + // Get the Arrow schema from the table + struct ArrowSchema* c_schema_ptr = nullptr; + char* error_message = nullptr; + if (const LanceDBError result = lancedb_table_arrow_schema( + table, + reinterpret_cast(&c_schema_ptr), + &error_message); result != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to get schema for index: " << index_name + << ". error: " << error_message << dendl; + lancedb_free_string(error_message); + return lancedb_error_to_errno(result); + } + + // Import the schema to Arrow C++ + const auto schema = arrow::ImportSchema(c_schema_ptr); + if (!schema.ok()) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to import schema for index: " << index_name + << ". error: " << schema.status().ToString() << dendl; + lancedb_free_arrow_schema(reinterpret_cast(c_schema_ptr)); + return -EINVAL; + } + + // Extract dimension from the "data" field + auto data_field = schema->get()->GetFieldByName(data_field_str); + if (!data_field) { + ldpp_dout(dpp, 1) << "ERROR: s3vector schema missing " << data_field_str << " field for index: " << index_name << dendl; + lancedb_free_arrow_schema(reinterpret_cast(c_schema_ptr)); + return -EINVAL; + } + + if (data_field->type()->id() != arrow::Type::FIXED_SIZE_LIST) { + ldpp_dout(dpp, 1) << "ERROR: s3vector " << data_field_str << " field is not a FixedSizeList for index: " << index_name << dendl; + lancedb_free_arrow_schema(reinterpret_cast(c_schema_ptr)); + return -EINVAL; + } + + auto fixed_size_list_type = std::static_pointer_cast(data_field->type()); + dimension = fixed_size_list_type->list_size(); + lancedb_free_arrow_schema(reinterpret_cast(c_schema_ptr)); return 0; -} + } -void vector_item_t::dump(ceph::Formatter* f) const { - f->open_object_section(""); - ::encode_json("key", key, f); - rgw::s3vector::encode_json("data", data, f); - ::encode_json("metadata", metadata, f); - f->close_section(); -} + int get_index(const get_index_t& configuration, const std::string& region, const std::string& account, DoutPrefixProvider* dpp, optional_yield y, get_index_reply_t& reply) { + log_configuration(dpp, "GetIndex", configuration); + LanceDBConnection* conn = connect(dpp, configuration.vector_bucket_name); + if (!conn) { + return -EIO; + } + LanceDBTable* table = lancedb_connection_open_table(conn, configuration.index_name.c_str()); + if (table == nullptr) { + lancedb_connection_free(conn); + return -ENOENT; + } + + if (int ret = get_vector_dimension(configuration.index_name, table, dpp, reply.dimension); ret < 0) { + lancedb_connection_free(conn); + lancedb_table_free(table); + return ret; + } + + reply.data_type = "float32"; + reply.distance_metric = get_table_distance_metric(table, dpp); + reply.index_name = configuration.index_name; + reply.vector_bucket_name = configuration.vector_bucket_name; + + if (configuration.index_arn) { + reply.index_arn = configuration.index_arn->to_string(); + } else { + reply.index_arn = index_arn( + region, + account, + configuration.vector_bucket_name, + configuration.index_name).to_string(); + } -void vector_item_t::decode_json(JSONObj* obj) { - JSONDecoder::decode_json("key", key, obj, true); - if (key.empty()) { - throw JSONDecoder::err("vector key must be specified"); + reply.creation_time = get_table_creation_time(table, dpp); + // reply.non_filterable_metadata_keys - empty for now, TODO: store and retrieve from table metadata + lancedb_table_free(table); + lancedb_connection_free(conn); + return 0; } - rgw::s3vector::decode_json("data", data, obj); - if (data.empty()) { - throw JSONDecoder::err("vector data cannot be empty"); + // list indexes + + void list_indexes_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + ::encode_json("maxResults", max_results, f); + ::encode_json("nextToken", next_token, f); + ::encode_json("prefix", prefix, f); + if (vector_bucket_arn) { + ::encode_json("vectorBucketArn", vector_bucket_arn->to_string(), f); + } else { + ::encode_json("vectorBucketName", vector_bucket_name, f); + } + f->close_section(); } - JSONDecoder::decode_json("metadata", metadata, obj); + void list_indexes_t::decode_json(JSONObj* obj) { + JSONDecoder::decode_json("maxResults", max_results, default_max_results, obj); + JSONDecoder::decode_json("nextToken", next_token, obj); + JSONDecoder::decode_json("prefix", prefix, obj); + decode_vector_bucket_name(vector_bucket_name, vector_bucket_arn, obj); -} + if (max_results < 1 || max_results > 500) { + throw JSONDecoder::err(fmt::format("maxResults must be between 1 and 500, got {}", max_results)); + } -void put_vectors_t::dump(ceph::Formatter* f) const { - f->open_object_section(""); - ::encode_json("indexArn", index_arn, f); - ::encode_json("indexName", index_name, f); - ::encode_json("vectorBucketName", vector_bucket_name, f); - f->open_array_section("vectors"); - for (const auto& vector : vectors) { - vector.dump(f); - } - f->close_section(); - f->close_section(); -} + if (!next_token.empty() && (next_token.length() < 1 || next_token.length() > 512)) { + throw JSONDecoder::err(fmt::format("nextToken length must be between 1 and 512, got {}", next_token.length())); + } -void put_vectors_t::decode_json(JSONObj* obj) { - decode_name_or_arn("indexName", "indexArn", index_name, index_arn, obj); - decode_name("vectorBucketName", vector_bucket_name, obj); - JSONDecoder::decode_json("vectors", vectors, obj, true); + if (!prefix.empty() && (prefix.length() < 1 || prefix.length() > 63)) { + throw JSONDecoder::err(fmt::format("prefix length must be between 1 and 63, got {}", prefix.length())); + } + } - if (vectors.empty() or vectors.size() > 500) { - throw JSONDecoder::err(fmt::format("vectors array must contain 1-500 items, got {}", vectors.size())); + void index_summary_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + ::encode_json("creationTime", creation_time, f); + ::encode_json("indexArn", index_arn, f); + ::encode_json("indexName", index_name, f); + ::encode_json("vectorBucketName", vector_bucket_name, f); + f->close_section(); } -} -int delete_vector_bucket_policy(const delete_vector_bucket_policy_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { - JSONFormatter f; - configuration.dump(&f); - std::stringstream ss; - f.flush(ss); - ldpp_dout(dpp, 20) << "INFO: executing s3vector DeleteVectorBucketPolicy with: " << ss.str() << dendl; + void index_summary_t::decode_json(JSONObj* obj) { + JSONDecoder::decode_json("creationTime", creation_time, obj); + JSONDecoder::decode_json("indexArn", index_arn, obj); + JSONDecoder::decode_json("indexName", index_name, obj); + JSONDecoder::decode_json("vectorBucketName", vector_bucket_name, obj); + } + + void list_indexes_reply_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + f->open_array_section("indexes"); + for (const auto& index : indexes) { + index.dump(f); + } + f->close_section(); + ::encode_json("nextToken", next_token, f); + f->close_section(); + } + + int list_indexes(const list_indexes_t& configuration, DoutPrefixProvider* dpp, optional_yield y, list_indexes_reply_t& reply) { + log_configuration(dpp, "ListIndexes", configuration); + LanceDBConnection* conn = connect(dpp, configuration.vector_bucket_name); + if (!conn) { + return -EIO; + } + LanceDBTableNamesBuilder* builder = lancedb_connection_table_names_builder(conn); + if (!builder) { + lancedb_connection_free(conn); + return -EIO; + } + if (builder = lancedb_table_names_builder_limit(builder, configuration.max_results); !builder) { + lancedb_connection_free(conn); + return -EIO; + } + if (builder = lancedb_table_names_builder_start_after(builder, configuration.next_token.c_str()); !builder) { + lancedb_connection_free(conn); + return -EIO; + } + char** table_names; + size_t name_count; + char* error_message; + if (const LanceDBError result = lancedb_table_names_builder_execute(builder, &table_names, &name_count, &error_message); result != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: s3vector listing indexes of: " << configuration.vector_bucket_name << ", error: " << error_message << dendl; + lancedb_free_string(error_message); + lancedb_connection_free(conn); + return lancedb_error_to_errno(result); + } + const bool has_prefix = configuration.prefix.empty(); + ldpp_dout(dpp, 20) << "INFO: s3vector listing: " << configuration.vector_bucket_name << ", found: " << name_count << " indexes" << dendl; + for (size_t i = 0; i < name_count; i++) { + // TODO: once/if this is resolved: https://github.com/lancedb/lancedb/issues/2895 we can prefix filtering at the builder level + std::string_view name{table_names[i]}; + if (has_prefix) { + if (!name.starts_with(configuration.prefix)) { + ldpp_dout(dpp, 20) << "INFO: on s3vector listing, index: " << name << " is is filterted out" << dendl; + continue; + } + } + ceph_assert(configuration.vector_bucket_arn); + uint64_t creation_time = 0; + LanceDBTable* table = lancedb_connection_open_table(conn, table_names[i]); + if (table) { + creation_time = get_table_creation_time(table, dpp); + lancedb_table_free(table); + } + reply.indexes.emplace_back( + creation_time, + index_arn( + configuration.vector_bucket_arn->region, + configuration.vector_bucket_arn->account, + configuration.vector_bucket_name, + name).to_string(), + std::string(name), + configuration.vector_bucket_name + ); + ldpp_dout(dpp, 20) << "INFO: on s3vector listing, index: " << name << " is added to list" << dendl; + } + if (name_count > 0) { + reply.next_token = table_names[name_count-1]; + } + lancedb_free_table_names(table_names, name_count); + lancedb_connection_free(conn); return 0; -} + } -void get_vectors_t::dump(ceph::Formatter* f) const { - f->open_object_section(""); - ::encode_json("indexArn", index_arn, f); - ::encode_json("indexName", index_name, f); - ::encode_json("keys", keys, f); - ::encode_json("returnData", return_data, f); - ::encode_json("returnMetadata", return_metadata, f); - ::encode_json("vectorBucketName", vector_bucket_name, f); - f->close_section(); -} + // vector bucket operations: create, delete, get, list, put/get/delete policy + ////////////////////////////////////////////////////////////////////////////// + + // delete vector bucket -void get_vectors_t::decode_json(JSONObj* obj) { - decode_name_or_arn("indexName", "indexArn", index_name, index_arn, obj); - JSONDecoder::decode_json("keys", keys, obj, true); - JSONDecoder::decode_json("returnData", return_data, obj); - JSONDecoder::decode_json("returnMetadata", return_metadata, obj); - decode_name("vectorBucketName", vector_bucket_name, obj); + void delete_vector_bucket_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + if (vector_bucket_arn) { + ::encode_json("vectorBucketArn", vector_bucket_arn->to_string(), f); + } else { + ::encode_json("vectorBucketName", vector_bucket_name, f); + } + f->close_section(); + } - if (keys.empty() || keys.size() > 100) { - throw JSONDecoder::err(fmt::format("keys array must contain 1-100 items, got {}", keys.size())); + void delete_vector_bucket_t::decode_json(JSONObj* obj) { + decode_vector_bucket_name(vector_bucket_name, vector_bucket_arn, obj); } - for (const auto& key : keys) { - if (key.empty() || key.length() > 1024) { - throw JSONDecoder::err(fmt::format("each key must be 1-1024 characters long, got key of length {}", key.length())); + int delete_vector_bucket(const delete_vector_bucket_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { + log_configuration(dpp, "DeleteVectorBucket", configuration); + LanceDBConnection* conn = connect(dpp, configuration.vector_bucket_name); + if (!conn) { + return -EIO; + } + // force delete TODO: fail deletion if tables exists, unless a "force" flag is used + char* error_message; + if (LanceDBError err = lancedb_connection_drop_all_tables(conn, nullptr, &error_message); err != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: failed to drop content of s3vector bucket in: " << configuration.vector_bucket_name << ". error: " << error_message << dendl; + lancedb_free_string(error_message); + lancedb_connection_free(conn); + return -EIO; } + lancedb_connection_free(conn); + return 0; } -} -int put_vectors(const put_vectors_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { - JSONFormatter f; - configuration.dump(&f); - std::stringstream ss; - f.flush(ss); - ldpp_dout(dpp, 20) << "INFO: executing s3vector PutVectors with: " << ss.str() << dendl; + // delete vector bucket policy + + void delete_vector_bucket_policy_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + if (vector_bucket_arn) { + ::encode_json("vectorBucketArn", vector_bucket_arn->to_string(), f); + } else { + ::encode_json("vectorBucketName", vector_bucket_name, f); + } + f->close_section(); + } + + void delete_vector_bucket_policy_t::decode_json(JSONObj* obj) { + decode_vector_bucket_name(vector_bucket_name, vector_bucket_arn, obj); + } + + int delete_vector_bucket_policy(const delete_vector_bucket_policy_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { + log_configuration(dpp, "DeleteVectorBucketPolicy", configuration); + // policy TODO: implement return 0; -} + } -void list_vectors_t::dump(ceph::Formatter* f) const { - f->open_object_section(""); - ::encode_json("indexArn", index_arn, f); - ::encode_json("indexName", index_name, f); - ::encode_json("vectorBucketName", vector_bucket_name, f); - ::encode_json("maxResults", max_results, f); - ::encode_json("nextToken", next_token, f); - ::encode_json("returnData", return_data, f); - ::encode_json("returnMetadata", return_metadata, f); - if (segment_count > 0) { - ::encode_json("segmentCount", segment_count, f); - ::encode_json("segmentIndex", segment_index, f); - } - f->close_section(); -} + // create vector bucket -void list_vectors_t::decode_json(JSONObj* obj) { - decode_name_or_arn("indexName", "indexArn", index_name, index_arn, obj); - decode_name("vectorBucketName", vector_bucket_name, obj); - JSONDecoder::decode_json("maxResults", max_results, default_max_results, obj); - JSONDecoder::decode_json("nextToken", next_token, obj); - JSONDecoder::decode_json("returnData", return_data, obj); - JSONDecoder::decode_json("returnMetadata", return_metadata, obj); - JSONDecoder::decode_json("segmentCount", segment_count, obj); - JSONDecoder::decode_json("segmentIndex", segment_index, obj); + void create_vector_bucket_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + ::encode_json("vectorBucketName", vector_bucket_name, f); + f->close_section(); + } - if (max_results < 1 || max_results > 1000) { - throw JSONDecoder::err(fmt::format("maxResults must be between 1 and 1000, got {}", max_results)); + void create_vector_bucket_t::decode_json(JSONObj* obj) { + JSONDecoder::decode_json("vectorBucketName", vector_bucket_name, obj, true); + verify_name(vector_bucket_name); } - if (!next_token.empty() && (next_token.length() < 1 || next_token.length() > 2048)) { - throw JSONDecoder::err(fmt::format("nextToken length must be between 1 and 2048, got {}", next_token.length())); + void create_vector_bucket_reply_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + ::encode_json("vectorBucketArn", vector_bucket_arn, f); + f->close_section(); } - if (segment_count > 0) { - if (segment_count < 1 || segment_count > 16) { - throw JSONDecoder::err(fmt::format("segmentCount must be between 1 and 16, got {}", segment_count)); + int create_vector_bucket(const create_vector_bucket_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { + log_configuration(dpp, "CreateVectorBucket", configuration); + LanceDBConnection* conn = connect(dpp, configuration.vector_bucket_name); + if (!conn) { + return -EIO; } - if (segment_index >= segment_count) { - throw JSONDecoder::err(fmt::format("segmentIndex must be between 0 and segmentCount-1 ({}), got {}", segment_count - 1, segment_index)); + // verify connectivity by listing table names + char** table_names; + size_t name_count; + char* error_message; + if (const LanceDBError result = lancedb_connection_table_names(conn, &table_names, &name_count, &error_message); result != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to verify connectivity for: " << configuration.vector_bucket_name << ", error: " << error_message << dendl; + lancedb_free_string(error_message); + lancedb_connection_free(conn); + return lancedb_error_to_errno(result); } - } else if (segment_index > 0) { - throw JSONDecoder::err("segmentIndex requires segmentCount to be specified"); + lancedb_free_table_names(table_names, name_count); + lancedb_connection_free(conn); + return 0; } -} -int get_vectors(const get_vectors_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { - JSONFormatter f; - configuration.dump(&f); - std::stringstream ss; - f.flush(ss); - ldpp_dout(dpp, 20) << "INFO: executing s3vector GetVectors with: " << ss.str() << dendl; - return 0; -} + // get vector bucket -void list_vector_buckets_t::dump(ceph::Formatter* f) const { - f->open_object_section(""); - ::encode_json("maxResults", max_results, f); - ::encode_json("nextToken", next_token, f); - ::encode_json("prefix", prefix, f); - f->close_section(); -} + void get_vector_bucket_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + if (vector_bucket_arn) { + ::encode_json("vectorBucketArn", vector_bucket_arn->to_string(), f); + } else { + ::encode_json("vectorBucketName", vector_bucket_name, f); + } + f->close_section(); + } + + void get_vector_bucket_t::decode_json(JSONObj* obj) { + decode_vector_bucket_name(vector_bucket_name, vector_bucket_arn, obj); + } -void list_vector_buckets_t::decode_json(JSONObj* obj) { - JSONDecoder::decode_json("maxResults", max_results, default_max_results, obj); - JSONDecoder::decode_json("nextToken", next_token, obj); - JSONDecoder::decode_json("prefix", prefix, obj); + // list vector buckets - if (max_results < 1 || max_results > 500) { - throw JSONDecoder::err(fmt::format("maxResults must be between 1 and 500, got {}", max_results)); + void list_vector_buckets_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + ::encode_json("maxResults", max_results, f); + ::encode_json("nextToken", next_token, f); + ::encode_json("prefix", prefix, f); + f->close_section(); } - if (!next_token.empty() && (next_token.length() < 1 || next_token.length() > 512)) { - throw JSONDecoder::err(fmt::format("nextToken length must be between 1 and 512, got {}", next_token.length())); + void list_vector_buckets_t::decode_json(JSONObj* obj) { + JSONDecoder::decode_json("maxResults", max_results, default_max_results, obj); + JSONDecoder::decode_json("nextToken", next_token, obj); + JSONDecoder::decode_json("prefix", prefix, obj); + + if (max_results < 1 || max_results > 500) { + throw JSONDecoder::err(fmt::format("maxResults must be between 1 and 500, got {}", max_results)); + } + + if (!next_token.empty() && (next_token.length() < 1 || next_token.length() > 512)) { + throw JSONDecoder::err(fmt::format("nextToken length must be between 1 and 512, got {}", next_token.length())); + } + + if (!prefix.empty() && (prefix.length() < 1 || prefix.length() > 63)) { + throw JSONDecoder::err(fmt::format("prefix length must be between 1 and 63, got {}", prefix.length())); + } } - if (!prefix.empty() && (prefix.length() < 1 || prefix.length() > 63)) { - throw JSONDecoder::err(fmt::format("prefix length must be between 1 and 63, got {}", prefix.length())); + // put vector bucket policy + + void put_vector_bucket_policy_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + ::encode_json("policy", policy, f); + if (vector_bucket_arn) { + ::encode_json("vectorBucketArn", vector_bucket_arn->to_string(), f); + } else { + ::encode_json("vectorBucketName", vector_bucket_name, f); + } + f->close_section(); } -} -int list_vectors(const list_vectors_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { - JSONFormatter f; - configuration.dump(&f); - std::stringstream ss; - f.flush(ss); - ldpp_dout(dpp, 20) << "INFO: executing s3vector ListVectors with: " << ss.str() << dendl; - return 0; -} + void put_vector_bucket_policy_t::decode_json(JSONObj* obj) { + JSONDecoder::decode_json("policy", policy, obj, true); + decode_vector_bucket_name(vector_bucket_name, vector_bucket_arn, obj); -void get_vector_bucket_t::dump(ceph::Formatter* f) const { - f->open_object_section(""); - ::encode_json("vectorBucketArn", vector_bucket_arn, f); - ::encode_json("vectorBucketName", vector_bucket_name, f); - f->close_section(); -} + if (policy.empty()) { + throw JSONDecoder::err("policy must be specified and cannot be empty"); + } + // policy TODO: validate JSON + } -void get_vector_bucket_t::decode_json(JSONObj* obj) { - decode_name_or_arn("vectorBucketName", "vectorBucketArn", vector_bucket_name, vector_bucket_arn, obj); -} + // get vector bucket policy -void get_index_t::dump(ceph::Formatter* f) const { - f->open_object_section(""); - ::encode_json("indexArn", index_arn, f); - ::encode_json("indexName", index_name, f); - ::encode_json("vectorBucketName", vector_bucket_name, f); - f->close_section(); -} + void get_vector_bucket_policy_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + if (vector_bucket_arn) { + ::encode_json("vectorBucketArn", vector_bucket_arn->to_string(), f); + } else { + ::encode_json("vectorBucketName", vector_bucket_name, f); + } + f->close_section(); + } -void get_index_t::decode_json(JSONObj* obj) { - decode_name_or_arn("indexName", "indexArn", index_name, index_arn, obj); - decode_name("vectorBucketName", vector_bucket_name, obj); -} + void get_vector_bucket_policy_t::decode_json(JSONObj* obj) { + decode_vector_bucket_name(vector_bucket_name, vector_bucket_arn, obj); + } -void list_indexes_t::dump(ceph::Formatter* f) const { - f->open_object_section(""); - ::encode_json("maxResults", max_results, f); - ::encode_json("nextToken", next_token, f); - ::encode_json("prefix", prefix, f); - ::encode_json("vectorBucketArn", vector_bucket_arn, f); - ::encode_json("vectorBucketName", vector_bucket_name, f); - f->close_section(); -} + // vector operations: put/get/delete/query + ////////////////////////////////////////// + + // put vectors -void list_indexes_t::decode_json(JSONObj* obj) { - JSONDecoder::decode_json("maxResults", max_results, default_max_results, obj); - JSONDecoder::decode_json("nextToken", next_token, obj); - JSONDecoder::decode_json("prefix", prefix, obj); - decode_name_or_arn("vectorBucketName", "vectorBucketArn", vector_bucket_name, vector_bucket_arn, obj); + void vector_item_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + if (distance) { + f->dump_float("distance", *distance); + } + ::encode_json("key", key, f); + if (data) { + rgw::s3vector::encode_json("data", *data, f); + } + if (!metadata.empty()) { + ::encode_json("metadata", metadata, f); + } + f->close_section(); + } - if (max_results < 1 || max_results > 500) { - throw JSONDecoder::err(fmt::format("maxResults must be between 1 and 500, got {}", max_results)); + void vector_item_t::decode_json(JSONObj* obj) { + JSONDecoder::decode_json("key", key, obj, true); + if (key.empty()) { + throw JSONDecoder::err("vector key must be specified"); + } + data.emplace(); + rgw::s3vector::decode_json("data", *data, obj); + if (data->empty()) { + throw JSONDecoder::err("vector data cannot be empty"); + } + JSONDecoder::decode_json("metadata", metadata, obj); } - if (!next_token.empty() && (next_token.length() < 1 || next_token.length() > 512)) { - throw JSONDecoder::err(fmt::format("nextToken length must be between 1 and 512, got {}", next_token.length())); + void put_vectors_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + if (index_arn) { + ::encode_json("indexArn", index_arn->to_string(), f); + } else { + ::encode_json("indexName", index_name, f); + ::encode_json("vectorBucketName", vector_bucket_name, f); + } + f->open_array_section("vectors"); + for (const auto& vector : vectors) { + vector.dump(f); + } + f->close_section(); + f->close_section(); } - if (!prefix.empty() && (prefix.length() < 1 || prefix.length() > 63)) { - throw JSONDecoder::err(fmt::format("prefix length must be between 1 and 63, got {}", prefix.length())); + void put_vectors_t::decode_json(JSONObj* obj) { + decode_index_name(vector_bucket_name, index_name, obj); + JSONDecoder::decode_json("vectors", vectors, obj, true); + + if (vectors.empty() or vectors.size() > 500) { + throw JSONDecoder::err(fmt::format("vectors array must contain 1-500 items, got {}", vectors.size())); + } } -} -int get_index(const get_index_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { - JSONFormatter f; - configuration.dump(&f); - std::stringstream ss; - f.flush(ss); - ldpp_dout(dpp, 20) << "INFO: executing s3vector GetIndex with: " << ss.str() << dendl; + int put_vectors(const put_vectors_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { + log_configuration(dpp, "PutVectors", configuration); + LanceDBTable* table = open_table(dpp, configuration.vector_bucket_name, configuration.index_name); + if (!table) { + return -EIO; + } + if (configuration.vectors.empty()) { + ldpp_dout(dpp, 10) << "WARNING: s3vector no vectors provided" << dendl; + lancedb_table_free(table); + return 0; + } + + // get the schema and dimension from the table + unsigned int dimension = 0; + if (int ret = get_vector_dimension(configuration.index_name, table, dpp, dimension); ret < 0) { + lancedb_table_free(table); + return ret; + } + + struct ArrowSchema* c_schema_ptr = nullptr; + char* schema_error_message = nullptr; + if (const LanceDBError result = lancedb_table_arrow_schema( + table, + reinterpret_cast(&c_schema_ptr), + &schema_error_message); result != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to get schema for index: " << configuration.index_name + << ". error: " << schema_error_message << dendl; + lancedb_free_string(schema_error_message); + lancedb_table_free(table); + return lancedb_error_to_errno(result); + } + + const auto schema = arrow::ImportSchema(c_schema_ptr); + if (!schema.ok()) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to import schema for index: " << configuration.index_name + << ". error: " << schema.status().ToString() << dendl; + lancedb_free_arrow_schema(reinterpret_cast(c_schema_ptr)); + lancedb_table_free(table); + return -EINVAL; + } + + arrow::StringBuilder key_builder; + arrow::FloatBuilder float_builder; + arrow::FixedSizeListBuilder data_builder(arrow::default_memory_pool(), + std::make_unique(), + dimension); + // metadata TODO: metadata configuration should also be taken from the index configuration + unsigned int num_rows = 0; + for (const auto& vector : configuration.vectors) { + if (!vector.data) { + ldpp_dout(dpp, 5) << "WARNING: s3vector skipping vector with no data" << dendl; + continue; + } + if (vector.data->size() != dimension) { + ldpp_dout(dpp, 5) << "WARNING: s3vector vector dimension mismatch, expected " + << dimension << " got " << vector.data->size() << + ". skip vector with key: " << vector.key << dendl; + continue; + } + if (vector.key.empty()) { + ldpp_dout(dpp, 5) << "WARNING: s3vector skipping vector with empty key" << dendl; + continue; + } + // metadata TODO: check if metadata is allowed based on index config + // key column + key_builder.Append(vector.key).ok(); + // data column + auto list_builder = static_cast(data_builder.value_builder()); + for (const auto & value : vector.data.value()) { + list_builder->Append(value).ok(); + } + data_builder.Append().ok(); + ++num_rows; + } + + if (num_rows == 0) { + ldpp_dout(dpp, 1) << "ERROR: s3vector no valid vectors to insert" << dendl; + lancedb_table_free(table); + return -EINVAL; + } + + std::shared_ptr key_array, data_array; + key_builder.Finish(&key_array).ok(); + data_builder.Finish(&data_array).ok(); + + auto record_batch = arrow::RecordBatch::Make(*schema, num_rows, {key_array, data_array}); + ldpp_dout(dpp, 20) << "INFO: s3vector created record batch with " << num_rows << " rows" << dendl; + struct ArrowArray c_array; + struct ArrowSchema c_schema; + if (const auto status = arrow::ExportRecordBatch(*record_batch, &c_array, &c_schema); !status.ok()) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to export record batch to C ABI: " << status.ToString() << dendl; + if (c_schema.release) c_schema.release(&c_schema); + lancedb_table_free(table); + return -EINVAL; + } + + LanceDBRecordBatchReader* batch_reader = nullptr; + char* reader_error_message = nullptr; + if (const auto reader_result = lancedb_record_batch_reader_from_arrow( + reinterpret_cast(&c_array), + reinterpret_cast(&c_schema), + &batch_reader, + &reader_error_message); reader_result != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to create record batch reader from arrow arrays: " + << (reader_error_message ? reader_error_message : "unknown") << dendl; + lancedb_free_string(reader_error_message); + if (c_array.release) c_array.release(&c_array); + if (c_schema.release) c_schema.release(&c_schema); + lancedb_table_free(table); + return -EINVAL; + } + + // upsert data into table: update existing keys, insert new ones + ldpp_dout(dpp, 10) << "INFO: s3vector attempting to upsert " << num_rows << " vectors with dimension " << dimension << dendl; + const LanceDBMergeInsertConfig merge_config = { + .when_matched_update_all = 1, + .when_not_matched_insert_all = 1 + }; + char* error_message; + const LanceDBError result = lancedb_table_merge_insert( + table, + batch_reader, + key_columns, + num_key_columns, + &merge_config, + &error_message); + + if (result != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to upsert record batch to index. error: " << error_message << dendl; + lancedb_free_string(error_message); + lancedb_table_free(table); + return lancedb_error_to_errno(result); + } + lancedb_table_free(table); return 0; -} + } -void put_vector_bucket_policy_t::dump(ceph::Formatter* f) const { - f->open_object_section(""); - ::encode_json("policy", policy, f); - ::encode_json("vectorBucketArn", vector_bucket_arn, f); - ::encode_json("vectorBucketName", vector_bucket_name, f); - f->close_section(); -} + // get vectors -void put_vector_bucket_policy_t::decode_json(JSONObj* obj) { - JSONDecoder::decode_json("policy", policy, obj, true); - decode_name_or_arn("vectorBucketName", "vectorBucketArn", vector_bucket_name, vector_bucket_arn, obj); + void get_vectors_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + if (index_arn) { + ::encode_json("indexArn", index_arn->to_string(), f); + } else { + ::encode_json("indexName", index_name, f); + ::encode_json("vectorBucketName", vector_bucket_name, f); + } + ::encode_json("keys", keys, f); + ::encode_json("returnData", return_data, f); + ::encode_json("returnMetadata", return_metadata, f); + f->close_section(); + } + + void get_vectors_t::decode_json(JSONObj* obj) { + decode_index_name(vector_bucket_name, index_name, obj); + JSONDecoder::decode_json("keys", keys, obj, true); + JSONDecoder::decode_json("returnData", return_data, obj); + JSONDecoder::decode_json("returnMetadata", return_metadata, obj); + + if (keys.empty() || keys.size() > 100) { + throw JSONDecoder::err(fmt::format("keys array must contain 1-100 items, got {}", keys.size())); + } - if (policy.empty()) { - throw JSONDecoder::err("policy must be specified and cannot be empty"); + for (const auto& key : keys) { + if (key.empty() || key.length() > 1024) { + throw JSONDecoder::err(fmt::format("each key must be 1-1024 characters long, got key of length {}", key.length())); + } + } } - // TODO: validate JSON policy -} -int list_indexes(const list_indexes_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { - JSONFormatter f; - configuration.dump(&f); - std::stringstream ss; - f.flush(ss); - ldpp_dout(dpp, 20) << "INFO: executing s3vector ListIndexes with: " << ss.str() << dendl; + void get_vectors_reply_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + f->open_array_section("vectors"); + for (const auto& vector : vectors) { + vector.dump(f); + } + f->close_section(); + f->close_section(); + } + + int populate_vectors_from_arrow( + DoutPrefixProvider* dpp, + struct ArrowArray** c_arrays_ptr, + struct ArrowSchema* c_schema_ptr, + std::vector& vectors, + const std::string& index_name, + bool use_data, + bool use_distance, + bool vector_query) { + unsigned long num_columns = 1; + if (use_data) ++num_columns; + if (vector_query) ++num_columns; + if (auto schema = arrow::ImportSchema(c_schema_ptr); schema.ok()) { + if (auto array = arrow::ImportRecordBatch(reinterpret_cast(*c_arrays_ptr), *schema); array.ok()) { + const auto& record_batch = *array; + // return by rows instead of columns + for (auto row = 0U; row < record_batch->num_rows(); row++) { + const auto record_num_columns = static_cast(record_batch->num_columns()); + if (record_num_columns != num_columns) { + ldpp_dout(dpp, 1) << "ERROR: s3vector got invalid number of columns in record batch for index: " << + index_name << ". got: " << record_num_columns << " expected: " << num_columns << dendl; + lancedb_free_arrow_schema(reinterpret_cast(c_schema_ptr)); + return -EINVAL; + } + vector_item_t vector_item; + if (use_data) vector_item.data.emplace(); + for (auto col = 0U; col < num_columns; col++) { + auto column = record_batch->column(col); + auto field = record_batch->schema()->field(col); + + if (field->name() == key_field_str) { + const auto key_array = std::static_pointer_cast(column); + if (!key_array->IsNull(row)) { + vector_item.key = key_array->GetString(row); + } else { + vector_item.key = ""; + } + } else if (field->name() == data_field_str && vector_item.data) { + const auto data_array = std::static_pointer_cast(column); + if (!data_array->IsNull(row)) { + const auto values = std::static_pointer_cast(data_array->values()); + const auto start = data_array->value_offset(row); + const auto length = data_array->value_length(); + for (auto i = 0; i < length; i++) { + vector_item.data->push_back(values->Value(start + i)); + } + } else { + ldpp_dout(dpp, 5) << "WARNING: s3vector got no data in record batch for index: " << index_name <name() == distance_field_str) { + if (!use_distance) continue; + const auto distance_array = std::static_pointer_cast(column); + if (!distance_array->IsNull(row)) { + vector_item.distance = distance_array->Value(row); + } else { + ldpp_dout(dpp, 5) << "WARNING: s3vector got no distance in record batch for index: " << index_name <name() << + " in record batch for index: " << index_name <(c_schema_ptr)); + return -EINVAL; + } + } else { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to import schema from arrow C ABI for index: " << + index_name << ". error: " << schema.status().ToString() << dendl; + return -EINVAL; + } + + lancedb_free_arrow_arrays(reinterpret_cast(c_arrays_ptr), 1); + lancedb_free_arrow_schema(reinterpret_cast(c_schema_ptr)); return 0; -} + } -void get_vector_bucket_policy_t::dump(ceph::Formatter* f) const { - f->open_object_section(""); - ::encode_json("vectorBucketArn", vector_bucket_arn, f); - ::encode_json("vectorBucketName", vector_bucket_name, f); - f->close_section(); -} + int populate_vectors_from_query( + DoutPrefixProvider* dpp, + LanceDBQueryResult* query_result, + std::vector& vectors, + const std::string& index_name, + bool use_data, + bool use_distance, + bool vector_query) { + // distance can be used only with vector queries + ceph_assert(!use_distance || vector_query); + struct ArrowArray** c_arrays_ptr = nullptr; + struct ArrowSchema* c_schema_ptr = nullptr; + size_t count_out; + char* error_message; + if (const LanceDBError result = lancedb_query_result_to_arrow( + query_result, + reinterpret_cast(&c_arrays_ptr), + reinterpret_cast(&c_schema_ptr), + &count_out, + &error_message); result != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to convert query result to arrow arrays for index: " << index_name << ". error: " << error_message << dendl; + lancedb_free_string(error_message); + return lancedb_error_to_errno(result); + } + if (count_out == 0) { + // no results + lancedb_free_arrow_arrays(reinterpret_cast(c_arrays_ptr), 1); + lancedb_free_arrow_schema(reinterpret_cast(c_schema_ptr)); + return 0; + } + return populate_vectors_from_arrow(dpp, c_arrays_ptr, c_schema_ptr, vectors, index_name, use_data, use_distance, vector_query); + } -void get_vector_bucket_policy_t::decode_json(JSONObj* obj) { - decode_name_or_arn("vectorBucketName", "vectorBucketArn", vector_bucket_name, vector_bucket_arn, obj); -} + int get_vectors(const get_vectors_t& configuration, DoutPrefixProvider* dpp, optional_yield y, get_vectors_reply_t& reply) { + log_configuration(dpp, "GetVectors", configuration); + LanceDBTable* table = open_table(dpp, configuration.vector_bucket_name, configuration.index_name); + if (!table) { + return -EIO; + } -int put_vector_bucket_policy(const put_vector_bucket_policy_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { - JSONFormatter f; - configuration.dump(&f); - std::stringstream ss; - f.flush(ss); - ldpp_dout(dpp, 20) << "INFO: executing s3vector PutVectorBucketPolicy with: " << ss.str() << dendl; - return 0; -} + LanceDBQuery* query = lancedb_query_new(table); + if (!query) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to create query for index: " << configuration.index_name << dendl; + lancedb_table_free(table); + return -EIO; + } -int get_vector_bucket_policy(const get_vector_bucket_policy_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { - JSONFormatter f; - configuration.dump(&f); - std::stringstream ss; - f.flush(ss); - ldpp_dout(dpp, 20) << "INFO: executing s3vector GetVectorBucketPolicy with: " << ss.str() << dendl; - return 0; -} + char* error_message; + const unsigned long num_columns = (configuration.return_data ? 2 : 1); + // metadata TODO: implement fetching metadata + if (const LanceDBError result = lancedb_query_select(query, table_columns, num_columns, &error_message) ; result != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to set select columns for query on index: " << configuration.index_name << ". error: " << error_message << dendl; + lancedb_free_string(error_message); + lancedb_query_free(query); + lancedb_table_free(table); + return lancedb_error_to_errno(result); + } -void delete_vectors_t::dump(ceph::Formatter* f) const { - f->open_object_section(""); - ::encode_json("indexArn", index_arn, f); - ::encode_json("indexName", index_name, f); - ::encode_json("keys", keys, f); - ::encode_json("vectorBucketName", vector_bucket_name, f); - f->close_section(); -} + // build where filter for keys + std::ostringstream oss; + const auto keys_size = configuration.keys.size(); + for (size_t i = 0; i < keys_size; ++i) { + oss << "key = \"" << configuration.keys[i] << "\""; + if (i < keys_size - 1) { + oss << " OR "; + } + } + + if (const LanceDBError result = lancedb_query_where_filter(query, oss.str().c_str(), &error_message); result != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to set where filter for query on index: " << configuration.index_name << ". error: " << error_message << dendl; + lancedb_free_string(error_message); + lancedb_query_free(query); + lancedb_table_free(table); + return lancedb_error_to_errno(result); + } -void delete_vectors_t::decode_json(JSONObj* obj) { - decode_name_or_arn("indexName", "indexArn", index_name, index_arn, obj); - JSONDecoder::decode_json("keys", keys, obj, true); - decode_name("vectorBucketName", vector_bucket_name, obj); + LanceDBQueryResult* query_result = lancedb_query_execute(query); + if (!query_result) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to execute query on index: " << configuration.index_name << dendl; + lancedb_query_free(query); + lancedb_table_free(table); + return -EIO; + } - if (keys.empty() || keys.size() > 500) { - throw JSONDecoder::err(fmt::format("keys array must contain 1-500 items, got {}", keys.size())); + auto ret = populate_vectors_from_query(dpp, query_result, reply.vectors, configuration.index_name, configuration.return_data, false, false); + lancedb_table_free(table); + return ret; } - for (const auto& key : keys) { - if (key.empty() || key.length() > 1024) { - throw JSONDecoder::err(fmt::format("each key must be 1-1024 characters long, got key of length {}", key.length())); + // list vectors + + void list_vectors_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + if (index_arn) { + ::encode_json("indexArn", index_arn->to_string(), f); + } else { + ::encode_json("indexName", index_name, f); + ::encode_json("vectorBucketName", vector_bucket_name, f); } + ::encode_json("maxResults", max_results, f); + ::encode_json("nextToken", offset, f); + ::encode_json("returnData", return_data, f); + ::encode_json("returnMetadata", return_metadata, f); + if (segment_count > 0) { + ::encode_json("segmentCount", segment_count, f); + ::encode_json("segmentIndex", segment_index, f); + } + f->close_section(); } -} -int delete_vectors(const delete_vectors_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { - JSONFormatter f; - configuration.dump(&f); - std::stringstream ss; - f.flush(ss); - ldpp_dout(dpp, 20) << "INFO: executing s3vector DeleteVectors with: " << ss.str() << dendl; - return 0; -} + void list_vectors_t::decode_json(JSONObj* obj) { + decode_index_name(vector_bucket_name, index_name, obj); + JSONDecoder::decode_json("maxResults", max_results, default_max_results, obj); + std::string next_token; + JSONDecoder::decode_json("nextToken", next_token, obj); + JSONDecoder::decode_json("returnData", return_data, obj); + JSONDecoder::decode_json("returnMetadata", return_metadata, obj); + JSONDecoder::decode_json("segmentCount", segment_count, obj); + JSONDecoder::decode_json("segmentIndex", segment_index, obj); + + if (max_results < 1 || max_results > 1000) { + throw JSONDecoder::err(fmt::format("maxResults must be between 1 and 1000, got {}", max_results)); + } -void query_vectors_t::dump(ceph::Formatter* f) const { - f->open_object_section(""); - if (!filter.empty()) { - ::encode_json("filter", filter, f); - } - ::encode_json("indexArn", index_arn, f); - ::encode_json("indexName", index_name, f); - rgw::s3vector::encode_json("queryVector", query_vector, f); - ::encode_json("returnDistance", return_distance, f); - ::encode_json("returnMetadata", return_metadata, f); - ::encode_json("topK", top_k, f); - ::encode_json("vectorBucketName", vector_bucket_name, f); - f->close_section(); -} + // according to the AWS spec the "nextToken" should be a string + // however, in lancedb, fetching all vectors in pages is done using offsets + /* if (!next_token.empty() && (next_token.length() < 1 || next_token.length() > 2048)) { + throw JSONDecoder::err(fmt::format("nextToken length must be between 1 and 2048, got {}", next_token.length())); + }*/ + + if (!next_token.empty()) { + const char* start = next_token.data(); + const char* end = start + next_token.length(); + const auto result = std::from_chars(start, end, offset); + + if (result.ec == std::errc::invalid_argument) { + throw JSONDecoder::err("failed to parse next token as offset"); + } + + if (result.ec == std::errc::result_out_of_range) { + throw JSONDecoder::err("out of range offset in next token"); + } + } -void query_vectors_t::decode_json(JSONObj* obj) { - JSONDecoder::decode_json("filter", filter, obj); - decode_name_or_arn("indexName", "indexArn", index_name, index_arn, obj); - rgw::s3vector::decode_json("queryVector", query_vector, obj); - JSONDecoder::decode_json("returnDistance", return_distance, obj); - JSONDecoder::decode_json("returnMetadata", return_metadata, obj); - JSONDecoder::decode_json("topK", top_k, obj, true); - decode_name("vectorBucketName", vector_bucket_name, obj); + if (segment_count > 0) { + if (segment_count < 1 || segment_count > 16) { + throw JSONDecoder::err(fmt::format("segmentCount must be between 1 and 16, got {}", segment_count)); + } + if (segment_index >= segment_count) { + throw JSONDecoder::err(fmt::format("segmentIndex must be between 0 and segmentCount-1 ({}), got {}", segment_count - 1, segment_index)); + } + } else if (segment_index > 0) { + throw JSONDecoder::err("segmentIndex requires segmentCount to be specified"); + } + } - if (top_k < 1) { - throw JSONDecoder::err(fmt::format("topK must be at least 1, got {}", top_k)); + void list_vectors_reply_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + ::encode_json("nextToken", next_token, f); + f->open_array_section("vectors"); + for (const auto& vector : vectors) { + vector.dump(f); + } + f->close_section(); + f->close_section(); } - if (query_vector.empty()) { - throw JSONDecoder::err("queryVector cannot be empty"); + int list_vectors(const list_vectors_t& configuration, DoutPrefixProvider* dpp, optional_yield y, list_vectors_reply_t& reply) { + log_configuration(dpp, "ListVectors", configuration); + LanceDBTable* table = open_table(dpp, configuration.vector_bucket_name, configuration.index_name); + if (!table) { + return -EIO; + } + + LanceDBQuery* query = lancedb_query_new(table); + if (!query) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to create query for index: " << configuration.index_name << dendl; + lancedb_table_free(table); + return -EIO; + } + + char* error_message; + const unsigned long num_columns = (configuration.return_data ? 2 : 1); + // metadata TODO: implement metadata based queries + if (const LanceDBError result = lancedb_query_select(query, table_columns, num_columns, &error_message) ; result != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to set select columns for query on index: " << configuration.index_name << ". error: " << error_message << dendl; + lancedb_free_string(error_message); + lancedb_query_free(query); + lancedb_table_free(table); + return lancedb_error_to_errno(result); + } + + if (const LanceDBError result = lancedb_query_limit(query, configuration.max_results, &error_message) ; result != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to set query limit on index: " << configuration.index_name << ". error: " << error_message << dendl; + lancedb_free_string(error_message); + lancedb_query_free(query); + lancedb_table_free(table); + return lancedb_error_to_errno(result); + } + + if (configuration.offset > 0) { + if (const LanceDBError result = lancedb_query_offset(query, configuration.offset, &error_message) ; result != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to set query offset on index: " << configuration.index_name << ". error: " << error_message << dendl; + lancedb_free_string(error_message); + lancedb_query_free(query); + lancedb_table_free(table); + return lancedb_error_to_errno(result); + } + } + + LanceDBQueryResult* query_result = lancedb_query_execute(query); + if (!query_result) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to execute query on index: " << configuration.index_name << dendl; + lancedb_query_free(query); + lancedb_table_free(table); + return -EIO; + } + + int ret; + if (ret = populate_vectors_from_query(dpp, query_result, reply.vectors, configuration.index_name, configuration.return_data, false, false); ret == 0) { + const auto total_row_count = lancedb_table_count_rows(table); + const auto next_offset = reply.vectors.size() + configuration.offset; + if (next_offset < total_row_count) { + reply.next_token = std::to_string(next_offset); + } + } + lancedb_table_free(table); + return ret; } - // TODO: validate filter -} + // delete vectors -int query_vectors(const query_vectors_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { - JSONFormatter f; - configuration.dump(&f); - std::stringstream ss; - f.flush(ss); - ldpp_dout(dpp, 20) << "INFO: executing s3vector QueryVectors with: " << ss.str() << dendl; - return 0; -} + void delete_vectors_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + if (index_arn) { + ::encode_json("indexArn", index_arn->to_string(), f); + } else { + ::encode_json("indexName", index_name, f); + ::encode_json("vectorBucketName", vector_bucket_name, f); + } + ::encode_json("keys", keys, f); + f->close_section(); + } + + void delete_vectors_t::decode_json(JSONObj* obj) { + decode_index_name(vector_bucket_name, index_name, obj); + JSONDecoder::decode_json("keys", keys, obj, true); + + if (keys.empty() || keys.size() > 500) { + throw JSONDecoder::err(fmt::format("keys array must contain 1-500 items, got {}", keys.size())); + } + + for (const auto& key : keys) { + if (key.empty() || key.length() > 1024) { + throw JSONDecoder::err(fmt::format("each key must be 1-1024 characters long, got key of length {}", key.length())); + } + } + } + + int delete_vectors(const delete_vectors_t& configuration, DoutPrefixProvider* dpp, optional_yield y) { + log_configuration(dpp, "DeleteVectors", configuration); + LanceDBTable* table = open_table(dpp, configuration.vector_bucket_name, configuration.index_name); + if (!table) { + return -EIO; + } + // build where filter for keys + std::ostringstream oss; + const auto keys_size = configuration.keys.size(); + for (size_t i = 0; i < keys_size; ++i) { + oss << "key = \"" << configuration.keys[i] << "\""; + if (i < keys_size - 1) { + oss << " OR "; + } + } + char* error_message; + LanceDBError result = lancedb_table_delete(table, oss.str().c_str(), &error_message); + if (result != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to set select columns for query on index: " << configuration.index_name << ". error: " << error_message << dendl; + lancedb_free_string(error_message); + } + lancedb_table_free(table); + return lancedb_error_to_errno(result); + } + + // query vectors + + void query_vectors_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + if (!filter.empty()) { + ::encode_json("filter", filter, f); + } + if (index_arn) { + ::encode_json("indexArn", index_arn->to_string(), f); + } else { + ::encode_json("indexName", index_name, f); + ::encode_json("vectorBucketName", vector_bucket_name, f); + } + rgw::s3vector::encode_json("queryVector", query_vector, f); + ::encode_json("returnDistance", return_distance, f); + ::encode_json("returnMetadata", return_metadata, f); + ::encode_json("topK", top_k, f); + f->close_section(); + } + + void query_vectors_t::decode_json(JSONObj* obj) { + JSONDecoder::decode_json("filter", filter, obj); + decode_index_name(vector_bucket_name, index_name, obj); + rgw::s3vector::decode_json("queryVector", query_vector, obj); + JSONDecoder::decode_json("returnDistance", return_distance, obj); + JSONDecoder::decode_json("returnMetadata", return_metadata, obj); + JSONDecoder::decode_json("topK", top_k, obj, true); + + if (top_k < 1) { + throw JSONDecoder::err(fmt::format("topK must be at least 1, got {}", top_k)); + } + + if (query_vector.empty()) { + throw JSONDecoder::err("queryVector cannot be empty"); + } + + // metadata TODO: validate filter + } + + void query_vectors_reply_t::dump(ceph::Formatter* f) const { + f->open_object_section(""); + ::encode_json("distanceMetric", distance_metric_to_string(distance_metric), f); + f->open_array_section("vectors"); + for (const auto& vector : vectors) { + vector.dump(f); + } + f->close_section(); + f->close_section(); + } + + int query_vectors(const query_vectors_t& configuration, DoutPrefixProvider* dpp, optional_yield y, query_vectors_reply_t& reply) { + log_configuration(dpp, "QueryVectors", configuration); + LanceDBTable* table = open_table(dpp, configuration.vector_bucket_name, configuration.index_name); + if (!table) { + return -EIO; + } + + unsigned int table_dimension; + if (int ret = get_vector_dimension(configuration.index_name, table, dpp, table_dimension); ret < 0) { + lancedb_table_free(table); + return ret; + } + + const auto query_dimension = configuration.query_vector.size(); + if (table_dimension != query_dimension) { + ldpp_dout(dpp, 1) << "ERROR: s3vector query vector dimension (" << query_dimension << ") does not match index: " << + configuration.index_name << " vector dimension (" << table_dimension << ")" << dendl; + return -EINVAL; + } + + LanceDBVectorQuery* query = lancedb_vector_query_new(table, configuration.query_vector.data(), query_dimension); + if (!query) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to create vector query for index: " << configuration.index_name << dendl; + lancedb_table_free(table); + return -EIO; + } + + char* error_message; + constexpr auto num_columns = 1; + // metadata TODO: support metadata + if (const LanceDBError result = lancedb_vector_query_select(query, key_columns, num_columns, &error_message) ; result != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to set select columns for vector query on index: " << configuration.index_name << ". error: " << error_message << dendl; + lancedb_free_string(error_message); + lancedb_vector_query_free(query); + lancedb_table_free(table); + return lancedb_error_to_errno(result); + } + + if (const LanceDBError result = lancedb_vector_query_column(query, data_field, &error_message) ; result != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to set select columns for vector query on index: " << configuration.index_name << ". error: " << error_message << dendl; + lancedb_free_string(error_message); + lancedb_vector_query_free(query); + lancedb_table_free(table); + return lancedb_error_to_errno(result); + } + + if (const LanceDBError result = lancedb_vector_query_limit(query, configuration.top_k, &error_message) ; result != LANCEDB_SUCCESS) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to set top-k for vector query on index: " << configuration.index_name << ". error: " << error_message << dendl; + lancedb_free_string(error_message); + lancedb_vector_query_free(query); + lancedb_table_free(table); + return lancedb_error_to_errno(result); + } + + LanceDBQueryResult* query_result = lancedb_vector_query_execute(query); + if (!query_result) { + ldpp_dout(dpp, 1) << "ERROR: s3vector failed to execute query on index: " << configuration.index_name << dendl; + lancedb_vector_query_free(query); + lancedb_table_free(table); + return -EIO; + } + + int ret = populate_vectors_from_query(dpp, query_result, reply.vectors, configuration.index_name, false, configuration.return_distance, true); + reply.distance_metric = get_table_distance_metric(table, dpp); + lancedb_table_free(table); + return ret; + } } diff --git a/src/rgw/rgw_s3vector.h b/src/rgw/rgw_s3vector.h index 95ca7a0f78c..b1f1c7c2fc1 100644 --- a/src/rgw/rgw_s3vector.h +++ b/src/rgw/rgw_s3vector.h @@ -6,6 +6,7 @@ #include #include #include "include/encoding.h" +#include "rgw_arn.h" #include "common/async/yield_context.h" namespace ceph { @@ -14,14 +15,14 @@ class Formatter; class JSONObj; class DoutPrefixProvider; -struct rgw_bucket; - namespace rgw::sal { class Driver; } namespace rgw::s3vector { + enum class DistanceMetric { + UNKNOWN, COSINE, EUCLIDEAN, }; @@ -45,25 +46,18 @@ struct create_index_t { DistanceMetric distance_metric; std::string index_name; std::vector non_filterable_metadata_keys; - std::string vector_bucket_arn; + boost::optional vector_bucket_arn; std::string vector_bucket_name; - void encode(bufferlist& bl) const { - ENCODE_START(1, 1, bl); - // TODO - ENCODE_FINISH(bl); - } + void dump(ceph::Formatter* f) const; + void decode_json(JSONObj* obj); +}; - void decode(bufferlist::const_iterator& bl) { - DECODE_START(1, bl); - // TODO - DECODE_FINISH(bl); - } +struct create_index_reply_t { + std::string index_arn; void dump(ceph::Formatter* f) const; - void decode_json(JSONObj* obj); }; -WRITE_CLASS_ENCODER(create_index_t) /* { @@ -73,22 +67,15 @@ WRITE_CLASS_ENCODER(create_index_t) struct create_vector_bucket_t { std::string vector_bucket_name; - void encode(bufferlist& bl) const { - ENCODE_START(1, 1, bl); - // TODO - ENCODE_FINISH(bl); - } + void dump(ceph::Formatter* f) const; + void decode_json(JSONObj* obj); +}; - void decode(bufferlist::const_iterator& bl) { - DECODE_START(1, bl); - // TODO - DECODE_FINISH(bl); - } +struct create_vector_bucket_reply_t { + std::string vector_bucket_arn; void dump(ceph::Formatter* f) const; - void decode_json(JSONObj* obj); }; -WRITE_CLASS_ENCODER(create_vector_bucket_t) /* { @@ -98,26 +85,13 @@ WRITE_CLASS_ENCODER(create_vector_bucket_t) } */ struct delete_index_t { - std::string index_arn; + boost::optional index_arn; std::string index_name; std::string vector_bucket_name; - void encode(bufferlist& bl) const { - ENCODE_START(1, 1, bl); - // TODO - ENCODE_FINISH(bl); - } - - void decode(bufferlist::const_iterator& bl) { - DECODE_START(1, bl); - // TODO - DECODE_FINISH(bl); - } - void dump(ceph::Formatter* f) const; void decode_json(JSONObj* obj); }; -WRITE_CLASS_ENCODER(delete_index_t) /* { @@ -126,25 +100,12 @@ WRITE_CLASS_ENCODER(delete_index_t) } */ struct delete_vector_bucket_t { - std::string vector_bucket_arn; + boost::optional vector_bucket_arn; std::string vector_bucket_name; - void encode(bufferlist& bl) const { - ENCODE_START(1, 1, bl); - // TODO - ENCODE_FINISH(bl); - } - - void decode(bufferlist::const_iterator& bl) { - DECODE_START(1, bl); - // TODO - DECODE_FINISH(bl); - } - void dump(ceph::Formatter* f) const; void decode_json(JSONObj* obj); }; -WRITE_CLASS_ENCODER(delete_vector_bucket_t) /* { @@ -153,56 +114,31 @@ WRITE_CLASS_ENCODER(delete_vector_bucket_t) } */ struct delete_vector_bucket_policy_t { - std::string vector_bucket_arn; + boost::optional vector_bucket_arn; std::string vector_bucket_name; - void encode(bufferlist& bl) const { - ENCODE_START(1, 1, bl); - // TODO - ENCODE_FINISH(bl); - } - - void decode(bufferlist::const_iterator& bl) { - DECODE_START(1, bl); - // TODO - DECODE_FINISH(bl); - } - void dump(ceph::Formatter* f) const; void decode_json(JSONObj* obj); }; -WRITE_CLASS_ENCODER(delete_vector_bucket_policy_t) - using VectorData = std::vector; /* { + "distance": "float", "key": "string", - "data": {"float32": [float]}, + "data": {"float32": ["float"]}, "metadata": {} } */ struct vector_item_t { + std::optional distance; std::string key; - VectorData data; + std::optional data; std::string metadata; // JSON string - void encode(bufferlist& bl) const { - ENCODE_START(1, 1, bl); - // TODO - ENCODE_FINISH(bl); - } - - void decode(bufferlist::const_iterator& bl) { - DECODE_START(1, bl); - // TODO - DECODE_FINISH(bl); - } - void dump(ceph::Formatter* f) const; void decode_json(JSONObj* obj); }; -WRITE_CLASS_ENCODER(vector_item_t) /* { @@ -213,62 +149,47 @@ WRITE_CLASS_ENCODER(vector_item_t) } */ struct put_vectors_t { - std::string index_arn; + boost::optional index_arn; std::string index_name; std::string vector_bucket_name; std::vector vectors; - void encode(bufferlist& bl) const { - ENCODE_START(1, 1, bl); - // TODO - ENCODE_FINISH(bl); - } - - void decode(bufferlist::const_iterator& bl) { - DECODE_START(1, bl); - // TODO - DECODE_FINISH(bl); - } - void dump(ceph::Formatter* f) const; void decode_json(JSONObj* obj); }; -WRITE_CLASS_ENCODER(put_vectors_t) /* { "indexArn": "string", "indexName": "string", + "vectorBucketName": "string" "keys": ["string"], "returnData": boolean, "returnMetadata": boolean, - "vectorBucketName": "string" } */ struct get_vectors_t { - std::string index_arn; + boost::optional index_arn; std::string index_name; + std::string vector_bucket_name; std::vector keys; bool return_data = false; bool return_metadata = false; - std::string vector_bucket_name; - void encode(bufferlist& bl) const { - ENCODE_START(1, 1, bl); - // TODO - ENCODE_FINISH(bl); - } + void dump(ceph::Formatter* f) const; + void decode_json(JSONObj* obj); +}; - void decode(bufferlist::const_iterator& bl) { - DECODE_START(1, bl); - // TODO - DECODE_FINISH(bl); +/* + { + "vectors": [vector_item_t] } +*/ +struct get_vectors_reply_t { + std::vector vectors; void dump(ceph::Formatter* f) const; - void decode_json(JSONObj* obj); }; -WRITE_CLASS_ENCODER(get_vectors_t) /* { @@ -284,33 +205,33 @@ WRITE_CLASS_ENCODER(get_vectors_t) } */ struct list_vectors_t { - std::string index_arn; + boost::optional index_arn; std::string index_name; std::string vector_bucket_name; static constexpr unsigned int default_max_results = 500; unsigned int max_results = default_max_results; - std::string next_token; + unsigned int offset = 0; // from nextToken bool return_data = false; bool return_metadata = false; unsigned int segment_count = 0; unsigned int segment_index = 0; - void encode(bufferlist& bl) const { - ENCODE_START(1, 1, bl); - // TODO - ENCODE_FINISH(bl); - } + void dump(ceph::Formatter* f) const; + void decode_json(JSONObj* obj); +}; - void decode(bufferlist::const_iterator& bl) { - DECODE_START(1, bl); - // TODO - DECODE_FINISH(bl); +/* + { + "nextToken": "string", + "vectors": [vector_item_t] } +*/ +struct list_vectors_reply_t { + std::string next_token; + std::vector vectors; void dump(ceph::Formatter* f) const; - void decode_json(JSONObj* obj); }; -WRITE_CLASS_ENCODER(list_vectors_t) /* { @@ -325,22 +246,9 @@ struct list_vector_buckets_t { std::string next_token; std::string prefix; - void encode(bufferlist& bl) const { - ENCODE_START(1, 1, bl); - // TODO - ENCODE_FINISH(bl); - } - - void decode(bufferlist::const_iterator& bl) { - DECODE_START(1, bl); - // TODO - DECODE_FINISH(bl); - } - void dump(ceph::Formatter* f) const; void decode_json(JSONObj* obj); }; -WRITE_CLASS_ENCODER(list_vector_buckets_t) /* { @@ -349,25 +257,12 @@ WRITE_CLASS_ENCODER(list_vector_buckets_t) } */ struct get_vector_bucket_t { - std::string vector_bucket_arn; + boost::optional vector_bucket_arn; std::string vector_bucket_name; - void encode(bufferlist& bl) const { - ENCODE_START(1, 1, bl); - // TODO - ENCODE_FINISH(bl); - } - - void decode(bufferlist::const_iterator& bl) { - DECODE_START(1, bl); - // TODO - DECODE_FINISH(bl); - } - void dump(ceph::Formatter* f) const; void decode_json(JSONObj* obj); }; -WRITE_CLASS_ENCODER(get_vector_bucket_t) /* { @@ -377,26 +272,42 @@ WRITE_CLASS_ENCODER(get_vector_bucket_t) } */ struct get_index_t { - std::string index_arn; + boost::optional index_arn; std::string index_name; std::string vector_bucket_name; - void encode(bufferlist& bl) const { - ENCODE_START(1, 1, bl); - // TODO - ENCODE_FINISH(bl); - } + void dump(ceph::Formatter* f) const; + void decode_json(JSONObj* obj); +}; - void decode(bufferlist::const_iterator& bl) { - DECODE_START(1, bl); - // TODO - DECODE_FINISH(bl); - } +/* + { + "index": { + "creationTime": number, + "dataType": "string", + "dimension": number, + "distanceMetric": "string", + "indexArn": "string", + "indexName": "string", + "metadataConfiguration": { + "nonFilterableMetadataKeys": [ "string" ] + }, + "vectorBucketName": "string" + } + } + */ +struct get_index_reply_t { + unsigned int creation_time; + std::string data_type; + unsigned int dimension; /* 1 - 4096 */ + DistanceMetric distance_metric; + std::string index_arn; + std::string index_name; + std::vector non_filterable_metadata_keys; + std::string vector_bucket_name; void dump(ceph::Formatter* f) const; - void decode_json(JSONObj* obj); }; -WRITE_CLASS_ENCODER(get_index_t) /* { @@ -412,25 +323,50 @@ struct list_indexes_t { unsigned int max_results = default_max_results; std::string next_token; std::string prefix; - std::string vector_bucket_arn; + boost::optional vector_bucket_arn; std::string vector_bucket_name; - void encode(bufferlist& bl) const { - ENCODE_START(1, 1, bl); - // TODO - ENCODE_FINISH(bl); - } + void dump(ceph::Formatter* f) const; + void decode_json(JSONObj* obj); +}; - void decode(bufferlist::const_iterator& bl) { - DECODE_START(1, bl); - // TODO - DECODE_FINISH(bl); - } +/* + { + "creationTime": number, + "indexArn": "string", + "indexName": "string", + "vectorBucketName": "string" + } + */ +struct index_summary_t { + unsigned int creation_time; + std::string index_arn; + std::string index_name; + std::string vector_bucket_name; void dump(ceph::Formatter* f) const; void decode_json(JSONObj* obj); }; -WRITE_CLASS_ENCODER(list_indexes_t) + +/* + { + "indexes": [ + { + "creationTime": number, + "indexArn": "string", + "indexName": "string", + "vectorBucketName": "string" + } + ], + "nextToken": "string" + } +*/ +struct list_indexes_reply_t { + std::vector indexes; + std::string next_token; + + void dump(ceph::Formatter* f) const; +}; /* { @@ -441,25 +377,12 @@ WRITE_CLASS_ENCODER(list_indexes_t) */ struct put_vector_bucket_policy_t { std::string policy; - std::string vector_bucket_arn; + boost::optional vector_bucket_arn; std::string vector_bucket_name; - void encode(bufferlist& bl) const { - ENCODE_START(1, 1, bl); - // TODO - ENCODE_FINISH(bl); - } - - void decode(bufferlist::const_iterator& bl) { - DECODE_START(1, bl); - // TODO - DECODE_FINISH(bl); - } - void dump(ceph::Formatter* f) const; void decode_json(JSONObj* obj); }; -WRITE_CLASS_ENCODER(put_vector_bucket_policy_t) /* { @@ -468,95 +391,88 @@ WRITE_CLASS_ENCODER(put_vector_bucket_policy_t) } */ struct get_vector_bucket_policy_t { - std::string vector_bucket_arn; + boost::optional vector_bucket_arn; std::string vector_bucket_name; - void encode(bufferlist& bl) const { - ENCODE_START(1, 1, bl); - // TODO - ENCODE_FINISH(bl); - } - - void decode(bufferlist::const_iterator& bl) { - DECODE_START(1, bl); - // TODO - DECODE_FINISH(bl); - } - void dump(ceph::Formatter* f) const; void decode_json(JSONObj* obj); }; -WRITE_CLASS_ENCODER(get_vector_bucket_policy_t) /* { "indexArn": "string", "indexName": "string", - "keys": ["string"], "vectorBucketName": "string" + "keys": ["string"], } */ struct delete_vectors_t { - std::string index_arn; + boost::optional index_arn; std::string index_name; - std::vector keys; std::string vector_bucket_name; - - void encode(bufferlist& bl) const { - ENCODE_START(1, 1, bl); - // TODO - ENCODE_FINISH(bl); - } - - void decode(bufferlist::const_iterator& bl) { - DECODE_START(1, bl); - // TODO - DECODE_FINISH(bl); - } + std::vector keys; void dump(ceph::Formatter* f) const; void decode_json(JSONObj* obj); }; -WRITE_CLASS_ENCODER(delete_vectors_t) /* { "filter": {}, "indexArn": "string", "indexName": "string", + "vectorBucketName": "string" "queryVector": {"float32": [float]}, "returnDistance": boolean, "returnMetadata": boolean, "topK": number, - "vectorBucketName": "string" } */ struct query_vectors_t { std::string filter; // JSON string - std::string index_arn; + boost::optional index_arn; std::string index_name; + std::string vector_bucket_name; VectorData query_vector; bool return_distance = false; bool return_metadata = false; unsigned int top_k; - std::string vector_bucket_name; - void encode(bufferlist& bl) const { - ENCODE_START(1, 1, bl); - // TODO - ENCODE_FINISH(bl); - } + void dump(ceph::Formatter* f) const; + void decode_json(JSONObj* obj); +}; - void decode(bufferlist::const_iterator& bl) { - DECODE_START(1, bl); - // TODO - DECODE_FINISH(bl); - } +/* +{ + "distanceMetric": "string", + "vectors": [vector_item_t] +} +*/ + +struct query_vectors_reply_t { + DistanceMetric distance_metric; + std::vector vectors; void dump(ceph::Formatter* f) const; - void decode_json(JSONObj* obj); }; -WRITE_CLASS_ENCODER(query_vectors_t) + +inline rgw::ARN index_arn(const std::string& zonegroup, const std::string& account, const std::string& vector_bucket_name, std::string_view index_name) { + return rgw::ARN(rgw::Partition::aws, + rgw::Service::s3vectors, + zonegroup, + account, + fmt::format("bucket/{}/index/{}", vector_bucket_name, index_name) + ); +} + +inline rgw::ARN vector_bucket_arn(const std::string& zonegroup, const std::string& account, const std::string& vector_bucket_name) { + return rgw::ARN(rgw::Partition::aws, + rgw::Service::s3vectors, + zonegroup, + account, + fmt::format("bucket/{}", vector_bucket_name) + ); +} int create_index(const create_index_t& configuration, DoutPrefixProvider* dpp, optional_yield y); int create_vector_bucket(const create_vector_bucket_t& configuration, DoutPrefixProvider* dpp, optional_yield y); @@ -564,14 +480,14 @@ int delete_index(const delete_index_t& configuration, DoutPrefixProvider* dpp, o int delete_vector_bucket(const delete_vector_bucket_t& configuration, DoutPrefixProvider* dpp, optional_yield y); int delete_vector_bucket_policy(const delete_vector_bucket_policy_t& configuration, DoutPrefixProvider* dpp, optional_yield y); int put_vectors(const put_vectors_t& configuration, DoutPrefixProvider* dpp, optional_yield y); -int get_vectors(const get_vectors_t& configuration, DoutPrefixProvider* dpp, optional_yield y); -int list_vectors(const list_vectors_t& configuration, DoutPrefixProvider* dpp, optional_yield y); -int get_index(const get_index_t& configuration, DoutPrefixProvider* dpp, optional_yield y); -int list_indexes(const list_indexes_t& configuration, DoutPrefixProvider* dpp, optional_yield y); +int get_vectors(const get_vectors_t& configuration, DoutPrefixProvider* dpp, optional_yield y, get_vectors_reply_t& reply); +int list_vectors(const list_vectors_t& configuration, DoutPrefixProvider* dpp, optional_yield y, list_vectors_reply_t& reply); +int get_index(const get_index_t& configuration, const std::string& region, const std::string& account, DoutPrefixProvider* dpp, optional_yield y, get_index_reply_t& reply); +int list_indexes(const list_indexes_t& configuration, DoutPrefixProvider* dpp, optional_yield y, list_indexes_reply_t& reply); int put_vector_bucket_policy(const put_vector_bucket_policy_t& configuration, DoutPrefixProvider* dpp, optional_yield y); int get_vector_bucket_policy(const get_vector_bucket_policy_t& configuration, DoutPrefixProvider* dpp, optional_yield y); int delete_vectors(const delete_vectors_t& configuration, DoutPrefixProvider* dpp, optional_yield y); -int query_vectors(const query_vectors_t& configuration, DoutPrefixProvider* dpp, optional_yield y); +int query_vectors(const query_vectors_t& configuration, DoutPrefixProvider* dpp, optional_yield y, query_vectors_reply_t& reply); } diff --git a/src/test/rgw/s3vectors/s3vector_test.py b/src/test/rgw/s3vectors/s3vector_test.py index e918a79aef3..97fef27557a 100644 --- a/src/test/rgw/s3vectors/s3vector_test.py +++ b/src/test/rgw/s3vectors/s3vector_test.py @@ -9,6 +9,8 @@ import subprocess import os import stat import string +from datetime import datetime, timedelta, timezone +from datetime import datetime, timezone import pytest import boto3 from botocore.config import Config @@ -148,20 +150,40 @@ def test_create_vector_bucket(): conn = connection() bucket_name = gen_bucket_name() result = conn.create_vector_bucket(vectorBucketName=bucket_name) + log.info('create_vector_bucket result: %s', result) assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + assert result['vectorBucketArn'] == 'arn:aws:s3vectors:::bucket/{}'.format(bucket_name) # cleanup _delete_all_vector_buckets(conn) +@pytest.mark.skip(reason="connection does not fail even with permission change") +@pytest.mark.vector_bucket_test +def test_create_vector_bucket_bad_path(): + conn = connection() + bucket_name = gen_bucket_name() + db_path = '/tmp/lancedb/' + os.makedirs(db_path, exist_ok=True) + os.chmod(db_path, 0o555) + try: + pytest.raises(conn.exceptions.ClientError, conn.create_vector_bucket, vectorBucketName=bucket_name) + finally: + os.chmod(db_path, 0o755) + + @pytest.mark.vector_bucket_test def test_get_vector_bucket(): conn = connection() bucket_name = gen_bucket_name() result = conn.create_vector_bucket(vectorBucketName=bucket_name) assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + bucket_arn = result['vectorBucketArn'] result = conn.get_vector_bucket(vectorBucketName=bucket_name) + log.info("get_vector_buckets result: %s", result) assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + result = conn.get_vector_bucket(vectorBucketArn=bucket_arn) log.info("get_vector_buckets result: %s", result) + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 invalid_name = bucket_name + '-invalid' pytest.raises(conn.exceptions.ClientError, conn.get_vector_bucket, vectorBucketName=invalid_name) # cleanup @@ -387,6 +409,13 @@ def test_create_index(): index_name = 'test-index' result = conn.create_index(vectorBucketName=bucket_name, indexName=index_name, dataType='float32', dimension=128, distanceMetric='euclidean') assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + assert result['indexArn'] == 'arn:aws:s3vectors:::bucket/{}/index/{}'.format(bucket_name, index_name) + # idempotent create with same definition should succeed + result = conn.create_index(vectorBucketName=bucket_name, indexName=index_name, dataType='float32', dimension=128, distanceMetric='euclidean') + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + # create with different dimension should fail + pytest.raises(conn.exceptions.ClientError, conn.create_index, + vectorBucketName=bucket_name, indexName=index_name, dataType='float32', dimension=64, distanceMetric='euclidean') # create an index on bucket that does not exist invalid_bucket_name = bucket_name + '-invalid' pytest.raises(conn.exceptions.ClientError, conn.create_index, vectorBucketName=invalid_bucket_name, indexName=index_name, dataType='float32', dimension=128, distanceMetric='euclidean') @@ -398,13 +427,23 @@ def test_create_index(): def test_get_index(): conn = connection() bucket_name = gen_bucket_name() + dimension = 128 result = conn.create_vector_bucket(vectorBucketName=bucket_name) assert result['ResponseMetadata']['HTTPStatusCode'] == 200 index_name = 'test-index' - result = conn.create_index(vectorBucketName=bucket_name, indexName=index_name, dataType='float32', dimension=128, distanceMetric='euclidean') + result = conn.create_index(vectorBucketName=bucket_name, indexName=index_name, dataType='float32', dimension=dimension, distanceMetric='euclidean') assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + index_arn = result['indexArn'] result = conn.get_index(vectorBucketName=bucket_name, indexName=index_name) assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + result = conn.get_index(indexArn=index_arn) + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + log.info('get_index result: %s', result) + assert result["index"]["dimension"] == dimension + assert result["index"]["indexName"] == index_name + assert result["index"]['indexArn'] == 'arn:aws:s3vectors:::bucket/{}/index/{}'.format(bucket_name, index_name) + assert result["index"]["creationTime"] > datetime.now(timezone.utc) - timedelta(days=1), "creationTime should be within the last day" + assert result["index"]["distanceMetric"] == "euclidean" # get an index from bucket that does not exist invalid_bucket_name = bucket_name + '-invalid' pytest.raises(conn.exceptions.ClientError, conn.get_index, vectorBucketName=invalid_bucket_name, indexName=index_name) @@ -430,7 +469,7 @@ def test_delete_index(): result = conn.delete_index(vectorBucketName=bucket_name, indexName=index_name) assert result['ResponseMetadata']['HTTPStatusCode'] == 200 # not implemented yet - #with pytest.raises(conn.exceptions.NoSuchIndex): + #with pytest.raises(conn.exceptions.ClientError): # result = conn.get_index(vectorBucketName=bucket_name, indexName=index_name) # cleanup _ = conn.delete_vector_bucket(vectorBucketName=bucket_name) @@ -450,19 +489,21 @@ def test_list_indexes(): assert result['ResponseMetadata']['HTTPStatusCode'] == 200 result = conn.list_indexes(vectorBucketName=bucket_name) assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + log.info('list_indexes result: %s', result) + index_names = [i['indexName'] for i in result['indexes']] + assert index_name1 in index_names + assert index_name2 in index_names + for idx in result['indexes']: + assert idx['creationTime'] > datetime.now(timezone.utc) - timedelta(days=1), "creationTime should be within the last day" # list indexs from bucket that does not exist invalid_bucket_name = bucket_name + '-invalid' pytest.raises(conn.exceptions.ClientError, conn.list_indexes, vectorBucketName=invalid_bucket_name) - # not implemented yet - #index_names = [i['IndexName'] for i in result['Indexes']] - #assert index_name1 in index_names - #assert index_name2 in index_names # cleanup _ = conn.delete_vector_bucket(vectorBucketName=bucket_name) -def generate_data(dimension): - return {'float32': [float(j) for j in range(dimension)]} +def generate_data(dimension, index=0): + return {'float32': [random.gauss(float(index), 1.0) for _ in range(dimension)]} def generate_vectors(num_vectors, dimension): @@ -470,11 +511,128 @@ def generate_vectors(num_vectors, dimension): for i in range(num_vectors): vectors.append({ 'key': 'vec-' + str(i), - 'data': generate_data(dimension) + 'data': generate_data(dimension, i) }) return vectors +def verify_get_vectors(conn, bucket_name, index_name, vector_ids, expected_dimension=None): + return_data = expected_dimension is not None + result = conn.get_vectors(vectorBucketName=bucket_name, indexName=index_name, + keys=vector_ids, returnData=return_data) + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + + assert 'vectors' in result + returned_vectors = result['vectors'] + num_expected = len(vector_ids) + assert len(returned_vectors) == num_expected, \ + f"expected {num_expected} vectors but got {len(returned_vectors)}" + + for vector in returned_vectors: + assert 'key' in vector, "vector should have a 'key' field" + assert vector['key'] in vector_ids, f"unexpected vector key: {vector['key']}" + + if return_data: + # verify data is present + assert 'data' in vector, \ + f"vector {vector['key']} should have 'data' field when returnData=True" + assert 'float32' in vector['data'], \ + f"vector {vector['key']} data should have 'float32' field" + + actual_dimension = len(vector['data']['float32']) + assert actual_dimension == expected_dimension, \ + f"vector {vector['key']} has dimension {actual_dimension}, expected {expected_dimension}" + else: + # verify data is NOT present + assert 'data' not in vector, \ + f"vector {vector['key']} should not have 'data' field when returnData=False" + + # verify all requested keys are present + returned_keys = [v['key'] for v in returned_vectors] + assert set(returned_keys) == set(vector_ids), \ + f"returned keys don't match requested keys. got {set(returned_keys)}, expected {set(vector_ids)}" + + log.info('get_vectors verification completed: %d vectors with returnData=%s', + len(returned_vectors), return_data) + + return returned_vectors + + +def verify_list_vectors_pagination(conn, bucket_name, index_name, expected_vectors, max_results, + expected_dimension=None): + return_data = expected_dimension is not None + total_vectors = len(expected_vectors) + all_retrieved_vectors = [] + next_token = None + page_count = 0 + expected_pages = (total_vectors + max_results - 1) // max_results # ceiling division + + while True: + page_count += 1 + if next_token: + result = conn.list_vectors(vectorBucketName=bucket_name, indexName=index_name, + maxResults=max_results, nextToken=next_token, returnData=return_data) + else: + result = conn.list_vectors(vectorBucketName=bucket_name, indexName=index_name, + maxResults=max_results, returnData=return_data) + + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + page_vectors = result.get('vectors', []) + + for vector in page_vectors: + assert 'key' in vector, "vector should have a 'key' field" + + if return_data: + # verify data is present + assert 'data' in vector, \ + f"vector {vector['key']} should have 'data' field when returnData=True" + assert 'float32' in vector['data'], \ + f"vector {vector['key']} data should have 'float32' field" + + actual_dimension = len(vector['data']['float32']) + assert actual_dimension == expected_dimension, \ + f"vector {vector['key']} has dimension {actual_dimension}, expected {expected_dimension}" + else: + # verify data is NOT present + assert 'data' not in vector, \ + f"vector {vector['key']} should not have 'data' field when returnData=False" + log.info('page %d returned %d vectors', page_count, len(page_vectors)) + all_retrieved_vectors.extend(page_vectors) + + if 'nextToken' in result and result['nextToken']: + # if there's a next page, this page should have max_results items + assert len(page_vectors) == max_results, \ + f"page {page_count} should have {max_results} vectors but has {len(page_vectors)}" + next_token = result['nextToken'] + else: + # last page - should have the remainder (or max_results if exact multiple) + expected_last_page_size = total_vectors % max_results + if expected_last_page_size == 0: + expected_last_page_size = max_results + assert len(page_vectors) == expected_last_page_size, \ + f"last page should have {expected_last_page_size} vectors but has {len(page_vectors)}" + break + + # verify we got the expected number of pages + assert page_count == expected_pages, \ + f"expected {expected_pages} pages but got {page_count}" + + assert len(all_retrieved_vectors) == len(expected_vectors) + expected_by_key = {v['key']: v for v in expected_vectors} + for retrieved in all_retrieved_vectors: + assert retrieved['key'] in expected_by_key, f"unexpected key: {retrieved['key']}" + if return_data: + expected = expected_by_key[retrieved['key']] + vector_pairs = zip(retrieved['data']['float32'], expected['data']['float32']) + assert all(abs(a - b) < 1e-6 for a, b in vector_pairs), \ + f"returned data don't match expected data for key {retrieved['key']}" + + log.info('pagination verification completed: %d vectors across %d pages', + len(all_retrieved_vectors), page_count) + + return all_retrieved_vectors, page_count + + @pytest.mark.vector_test def test_put_vectors(): conn = connection() @@ -491,23 +649,100 @@ def test_put_vectors(): _ = conn.delete_vector_bucket(vectorBucketName=bucket_name) +@pytest.mark.vector_test +def test_put_vectors_dimension_mismatch(): + conn = connection() + bucket_name = gen_bucket_name() + result = conn.create_vector_bucket(vectorBucketName=bucket_name) + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + index_name = 'test-index' + dimension = 128 + result = conn.create_index(vectorBucketName=bucket_name, indexName=index_name, dataType='float32', dimension=dimension, distanceMetric='euclidean') + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + # generate vectors with wrong dimension + wrong_dimension = 64 + vectors = generate_vectors(10, wrong_dimension) + # all vectors have wrong dimension, so put should fail + pytest.raises(conn.exceptions.ClientError, conn.put_vectors, + vectorBucketName=bucket_name, indexName=index_name, vectors=vectors) + # verify no vectors were inserted + result = conn.list_vectors(vectorBucketName=bucket_name, indexName=index_name, maxResults=100) + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + assert len(result.get('vectors', [])) == 0 + # mix of correct and wrong dimension vectors - only correct ones should be inserted + correct_vectors = generate_vectors(5, dimension) + wrong_vectors = generate_vectors(5, wrong_dimension) + # rename wrong vectors keys to avoid collisions + for i, v in enumerate(wrong_vectors): + v['key'] = f'wrong-{i}' + mixed_vectors = correct_vectors + wrong_vectors + result = conn.put_vectors(vectorBucketName=bucket_name, indexName=index_name, vectors=mixed_vectors) + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + # verify only the correct dimension vectors were inserted + result = conn.list_vectors(vectorBucketName=bucket_name, indexName=index_name, maxResults=100) + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + inserted_keys = [v['key'] for v in result.get('vectors', [])] + assert len(inserted_keys) == 5, f"expected 5 vectors but got {len(inserted_keys)}" + for i in range(5): + assert f'vec-{i}' in inserted_keys + for i in range(5): + assert f'wrong-{i}' not in inserted_keys + # cleanup + _ = conn.delete_vector_bucket(vectorBucketName=bucket_name) + + @pytest.mark.vector_test def test_get_vectors(): conn = connection() bucket_name = gen_bucket_name() + dimension = 128 result = conn.create_vector_bucket(vectorBucketName=bucket_name) assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + index_name = 'test-index' - result = conn.create_index(vectorBucketName=bucket_name, indexName=index_name, dataType='float32', dimension=128, distanceMetric='euclidean') + result = conn.create_index(vectorBucketName=bucket_name, indexName=index_name, + dataType='float32', dimension=dimension, distanceMetric='euclidean') assert result['ResponseMetadata']['HTTPStatusCode'] == 200 - vectors = generate_vectors(10, 128) + + num_vectors = 10 + vectors = generate_vectors(num_vectors, dimension) result = conn.put_vectors(vectorBucketName=bucket_name, indexName=index_name, vectors=vectors) assert result['ResponseMetadata']['HTTPStatusCode'] == 200 - vector_ids = ['vec-' + str(i) for i in range(10)] - result = conn.get_vectors(vectorBucketName=bucket_name, indexName=index_name, keys=vector_ids) + + # Get vectors with returnData=True + vector_ids = [f'vec-{i}' for i in range(num_vectors)] + returned_vectors = verify_get_vectors(conn, bucket_name, index_name, vector_ids, expected_dimension=dimension) + + log.info('test_get_vectors: successfully verified %d vectors with data', len(returned_vectors)) + + # cleanup + _ = conn.delete_vector_bucket(vectorBucketName=bucket_name) + + +@pytest.mark.vector_test +def test_get_vectors_without_data(): + conn = connection() + bucket_name = gen_bucket_name() + dimension = 128 + result = conn.create_vector_bucket(vectorBucketName=bucket_name) assert result['ResponseMetadata']['HTTPStatusCode'] == 200 - # not implemented yet - #assert len(result['Vectors']) == 10 + + index_name = 'test-index' + result = conn.create_index(vectorBucketName=bucket_name, indexName=index_name, + dataType='float32', dimension=dimension, distanceMetric='euclidean') + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + + num_vectors = 10 + vectors = generate_vectors(num_vectors, dimension) + result = conn.put_vectors(vectorBucketName=bucket_name, indexName=index_name, vectors=vectors) + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + + vector_ids = [f'vec-{i}' for i in range(num_vectors)] + returned_vectors = verify_get_vectors(conn, bucket_name, index_name, vector_ids) + + log.info('test_get_vectors_without_data: successfully verified %d vectors without data', + len(returned_vectors)) + # cleanup _ = conn.delete_vector_bucket(vectorBucketName=bucket_name) @@ -518,16 +753,151 @@ def test_list_vectors(): bucket_name = gen_bucket_name() result = conn.create_vector_bucket(vectorBucketName=bucket_name) assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + index_name = 'test-index' - result = conn.create_index(vectorBucketName=bucket_name, indexName=index_name, dataType='float32', dimension=128, distanceMetric='euclidean') + dimension = 8 + result = conn.create_index(vectorBucketName=bucket_name, indexName=index_name, + dataType='float32', dimension=dimension, distanceMetric='euclidean') assert result['ResponseMetadata']['HTTPStatusCode'] == 200 - vectors = generate_vectors(10, 128) + + total_vectors = 5 + max_results = 100 + vectors = generate_vectors(total_vectors, dimension) + result = conn.put_vectors(vectorBucketName=bucket_name, indexName=index_name, vectors=vectors) assert result['ResponseMetadata']['HTTPStatusCode'] == 200 - result = conn.list_vectors(vectorBucketName=bucket_name, indexName=index_name) + + all_retrieved_vectors, page_count = verify_list_vectors_pagination( + conn, bucket_name, index_name, vectors, max_results) + + assert page_count == 1, f"expected 1 pages but got {page_count}" + log.info('test_list_vectors: successfully verified %d vectors across %d pages', + len(all_retrieved_vectors), page_count) + + # cleanup + _ = conn.delete_vector_bucket(vectorBucketName=bucket_name) + + +@pytest.mark.vector_test +def test_list_vectors_with_data(): + """Test list_vectors with returnData=True to verify data is returned.""" + conn = connection() + bucket_name = gen_bucket_name() + result = conn.create_vector_bucket(vectorBucketName=bucket_name) + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + + index_name = 'test-index' + dimension = 128 + result = conn.create_index(vectorBucketName=bucket_name, indexName=index_name, + dataType='float32', dimension=dimension, distanceMetric='euclidean') + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + + num_vectors = 15 + vectors = generate_vectors(num_vectors, dimension) + result = conn.put_vectors(vectorBucketName=bucket_name, indexName=index_name, vectors=vectors) assert result['ResponseMetadata']['HTTPStatusCode'] == 200 - # not implemented yet - #assert len(result['Vectors']) == 10 + + all_retrieved_vectors, page_count = verify_list_vectors_pagination( + conn, bucket_name, index_name, vectors, max_results=100, expected_dimension=dimension) + + assert page_count == 1, f"expected 1 page but got {page_count}" + log.info('test_list_vectors_with_data: successfully verified %d vectors with data', + len(all_retrieved_vectors)) + + # cleanup + _ = conn.delete_vector_bucket(vectorBucketName=bucket_name) + + +@pytest.mark.vector_test +def test_list_vectors_without_data(): + """Test list_vectors with returnData=False to verify data is not returned.""" + conn = connection() + bucket_name = gen_bucket_name() + result = conn.create_vector_bucket(vectorBucketName=bucket_name) + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + + index_name = 'test-index' + dimension = 128 + result = conn.create_index(vectorBucketName=bucket_name, indexName=index_name, + dataType='float32', dimension=dimension, distanceMetric='euclidean') + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + + num_vectors = 15 + vectors = generate_vectors(num_vectors, dimension) + result = conn.put_vectors(vectorBucketName=bucket_name, indexName=index_name, vectors=vectors) + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + + all_retrieved_vectors, page_count = verify_list_vectors_pagination( + conn, bucket_name, index_name, vectors, max_results=100) + + assert page_count == 1, f"expected 1 page but got {page_count}" + log.info('test_list_vectors_without_data: successfully verified %d vectors without data', + len(all_retrieved_vectors)) + + # cleanup + _ = conn.delete_vector_bucket(vectorBucketName=bucket_name) + + +@pytest.mark.vector_test +def test_list_vectors_pagination(): + conn = connection() + bucket_name = gen_bucket_name() + result = conn.create_vector_bucket(vectorBucketName=bucket_name) + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + + index_name = 'test-index' + dimension = 128 + result = conn.create_index(vectorBucketName=bucket_name, indexName=index_name, + dataType='float32', dimension=dimension, distanceMetric='euclidean') + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + + # 27 vectors with page size 10 = 3 pages (10, 10, 7) + total_vectors = 27 + max_results = 10 + vectors = generate_vectors(total_vectors, dimension) + + result = conn.put_vectors(vectorBucketName=bucket_name, indexName=index_name, vectors=vectors) + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + + all_retrieved_vectors, page_count = verify_list_vectors_pagination( + conn, bucket_name, index_name, vectors, max_results, expected_dimension=dimension) + + assert page_count == 3, f"expected 3 pages but got {page_count}" + log.info('test_list_vectors: successfully verified %d vectors across %d pages', + len(all_retrieved_vectors), page_count) + + # cleanup + _ = conn.delete_vector_bucket(vectorBucketName=bucket_name) + + +@pytest.mark.vector_test +def test_list_vectors_exact_pagination(): + conn = connection() + bucket_name = gen_bucket_name() + result = conn.create_vector_bucket(vectorBucketName=bucket_name) + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + + index_name = 'test-index' + dimension = 128 + result = conn.create_index(vectorBucketName=bucket_name, indexName=index_name, + dataType='float32', dimension=dimension, distanceMetric='euclidean') + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + + # 30 vectors with page size 10 = 3 pages (10, 10, 10) - exact fit + total_vectors = 30 + max_results = 10 + vectors = generate_vectors(total_vectors, dimension) + + result = conn.put_vectors(vectorBucketName=bucket_name, indexName=index_name, vectors=vectors) + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + + all_retrieved_vectors, page_count = verify_list_vectors_pagination( + conn, bucket_name, index_name, vectors, max_results, expected_dimension=dimension) + + assert page_count == 3, f"expected 3 pages but got {page_count}" + log.info('test_list_vectors_exact_pagination: successfully verified %d vectors across %d pages', + len(all_retrieved_vectors), page_count) + # cleanup _ = conn.delete_vector_bucket(vectorBucketName=bucket_name) @@ -538,40 +908,95 @@ def test_delete_vectors(): bucket_name = gen_bucket_name() result = conn.create_vector_bucket(vectorBucketName=bucket_name) assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + index_name = 'test-index' - result = conn.create_index(vectorBucketName=bucket_name, indexName=index_name, dataType='float32', dimension=128, distanceMetric='euclidean') + dimension = 128 + result = conn.create_index(vectorBucketName=bucket_name, indexName=index_name, + dataType='float32', dimension=dimension, distanceMetric='euclidean') assert result['ResponseMetadata']['HTTPStatusCode'] == 200 - vectors = generate_vectors(10, 128) + + total_vectors = 20 + vectors = generate_vectors(total_vectors, dimension) result = conn.put_vectors(vectorBucketName=bucket_name, indexName=index_name, vectors=vectors) assert result['ResponseMetadata']['HTTPStatusCode'] == 200 - vector_ids = ['vec-' + str(i) for i in range(10)] - result = conn.delete_vectors(vectorBucketName=bucket_name, indexName=index_name, keys=vector_ids) + + # known vectors to delete: vec-2, vec-5, vec-7, vec-10, vec-12, vec-15, vec-17, vec-19 + # unknown vectors: vec-100, vec-999, nonexistent-key + vectors_to_delete = ['vec-2', 'vec-5', 'vec-7', 'vec-10', 'vec-12', 'vec-15', 'vec-17', 'vec-19', + 'vec-100', 'vec-999', 'nonexistent-key'] + result = conn.delete_vectors(vectorBucketName=bucket_name, indexName=index_name, keys=vectors_to_delete) assert result['ResponseMetadata']['HTTPStatusCode'] == 200 - result = conn.list_vectors(vectorBucketName=bucket_name, indexName=index_name) + + # list all vectors to verify deletion + result = conn.list_vectors(vectorBucketName=bucket_name, indexName=index_name, + maxResults=100, returnData=False) assert result['ResponseMetadata']['HTTPStatusCode'] == 200 - # not implemented yet - #assert len(result['Vectors']) == 0 + + remaining_vectors = result.get('vectors', []) + remaining_keys = [v['key'] for v in remaining_vectors] + expected_remaining_keys = [f'vec-{i}' for i in range(total_vectors) if f'vec-{i}' not in vectors_to_delete] + assert set(remaining_keys) == set(expected_remaining_keys), \ + f"remaining vector keys don't match expected. got {set(remaining_keys)}, expected {set(expected_remaining_keys)}" + # cleanup _ = conn.delete_vector_bucket(vectorBucketName=bucket_name) @pytest.mark.vector_test def test_query_vectors(): + dimension = 8 conn = connection() bucket_name = gen_bucket_name() result = conn.create_vector_bucket(vectorBucketName=bucket_name) assert result['ResponseMetadata']['HTTPStatusCode'] == 200 index_name = 'test-index' - result = conn.create_index(vectorBucketName=bucket_name, indexName=index_name, dataType='float32', dimension=128, distanceMetric='euclidean') + result = conn.create_index(vectorBucketName=bucket_name, indexName=index_name, dataType='float32', dimension=dimension, distanceMetric='euclidean') assert result['ResponseMetadata']['HTTPStatusCode'] == 200 - vectors = generate_vectors(10, 128) + vectors = generate_vectors(100, dimension) result = conn.put_vectors(vectorBucketName=bucket_name, indexName=index_name, vectors=vectors) assert result['ResponseMetadata']['HTTPStatusCode'] == 200 - query_vector = generate_data(128) - result = conn.query_vectors(vectorBucketName=bucket_name, indexName=index_name, queryVector=query_vector, topK=5) + top_k = 5 + for expected_index in [8, 17, 42, 99]: + query_vector = generate_data(dimension, expected_index) + result = conn.query_vectors(vectorBucketName=bucket_name, indexName=index_name, queryVector=query_vector, topK=top_k) + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + expected_key = 'vec-'+str(expected_index) + assert expected_key in [v['key'] for v in result['vectors']] + assert 'distance' not in [v for v in result['vectors']] + assert len(result['vectors']) == top_k + assert result['distanceMetric'] == 'euclidean' + log.info(result['vectors']) + + # cleanup + _ = conn.delete_vector_bucket(vectorBucketName=bucket_name) + + +@pytest.mark.vector_test +def test_query_vectors_with_distance(): + dimension = 8 + conn = connection() + bucket_name = gen_bucket_name() + result = conn.create_vector_bucket(vectorBucketName=bucket_name) + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + index_name = 'test-index' + result = conn.create_index(vectorBucketName=bucket_name, indexName=index_name, dataType='float32', dimension=dimension, distanceMetric='euclidean') + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + vectors = generate_vectors(100, dimension) + result = conn.put_vectors(vectorBucketName=bucket_name, indexName=index_name, vectors=vectors) assert result['ResponseMetadata']['HTTPStatusCode'] == 200 - # not implemented yet - #assert len(result['Results']) == 5 + top_k = 5 + for expected_index in [8, 17, 42, 99]: + query_vector = generate_data(dimension, expected_index) + result = conn.query_vectors(vectorBucketName=bucket_name, indexName=index_name, queryVector=query_vector, topK=top_k, returnDistance=True) + assert result['ResponseMetadata']['HTTPStatusCode'] == 200 + expected_key = 'vec-'+str(expected_index) + assert expected_key in [v['key'] for v in result['vectors']] + for v in result['vectors']: + assert 'distance' in v + assert len(result['vectors']) == top_k + assert result['distanceMetric'] == 'euclidean' + log.info(result['vectors']) + # cleanup _ = conn.delete_vector_bucket(vectorBucketName=bucket_name)