diff --git a/CMakeLists.txt b/CMakeLists.txt index 77ed0c34..91a39823 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -54,10 +54,6 @@ option(PAIMON_ENABLE_AVRO "Whether to enable avro file format" ON) option(PAIMON_ENABLE_ORC "Whether to enable orc file format" ON) option(PAIMON_ENABLE_JINDO "Whether to enable jindo file system" OFF) option(PAIMON_ENABLE_LUCENE "Whether to enable lucene index" OFF) -if(PAIMON_ENABLE_LUCENE) - message(FATAL_ERROR "PAIMON_ENABLE_LUCENE is temporarily disabled. Please configure with -DPAIMON_ENABLE_LUCENE=OFF." - ) -endif() if(PAIMON_ENABLE_ORC) add_definitions(-DPAIMON_ENABLE_ORC) @@ -302,12 +298,6 @@ endif() list(APPEND PAIMON_LINK_LIBS ${CMAKE_DL_LIBS}) list(APPEND PAIMON_SHARED_INSTALL_INTERFACE_LIBS ${CMAKE_DL_LIBS}) -if(PAIMON_ENABLE_LUMINA) - add_subdirectory(third_party/lumina EXCLUDE_FROM_ALL) - link_directories(third_party/lumina/lib) - install(FILES ${PROJECT_SOURCE_DIR}/third_party/lumina/lib/liblumina.so - DESTINATION ${CMAKE_INSTALL_LIBDIR}) -endif() if(PAIMON_ENABLE_LUCENE) set(PAIMON_DICT_DEST "share/paimon/dict") @@ -341,10 +331,6 @@ if(PAIMON_ENABLE_LANCE) include_directories("${CMAKE_SOURCE_DIR}/third_party/lance") endif() -if(PAIMON_ENABLE_LUMINA) - include_directories("${CMAKE_SOURCE_DIR}/third_party/lumina/include") -endif() - include_directories(SYSTEM ${ARROW_INCLUDE_DIR}) include_directories(SYSTEM ${TBB_INCLUDE_DIR}) diff --git a/ci/scripts/build_paimon.sh b/ci/scripts/build_paimon.sh index 8847caf4..235434a7 100755 --- a/ci/scripts/build_paimon.sh +++ b/ci/scripts/build_paimon.sh @@ -35,12 +35,18 @@ fi mkdir ${build_dir} pushd ${build_dir} +ENABLE_LUMINA="ON" +if [[ "${CC:-}" == *"gcc-8"* ]] || [[ "${CXX:-}" == *"g++-8"* ]]; then + ENABLE_LUMINA="OFF" +fi + CMAKE_ARGS=( "-G Ninja" "-DCMAKE_BUILD_TYPE=${build_type}" "-DPAIMON_BUILD_TESTS=ON" "-DPAIMON_ENABLE_JINDO=ON" - "-DPAIMON_ENABLE_LUCENE=OFF" + "-DPAIMON_ENABLE_LUMINA=${ENABLE_LUMINA}" + "-DPAIMON_ENABLE_LUCENE=ON" ) if [[ "${enable_sanitizer}" == "true" ]]; then diff --git a/cmake_modules/ThirdpartyToolchain.cmake b/cmake_modules/ThirdpartyToolchain.cmake index 669301bc..ec93aa50 100644 --- a/cmake_modules/ThirdpartyToolchain.cmake +++ b/cmake_modules/ThirdpartyToolchain.cmake @@ -281,6 +281,26 @@ else() endif() endif() +if(DEFINED ENV{PAIMON_LUMINA_URL}) + set(LUMINA_SOURCE_URL "$ENV{PAIMON_LUMINA_URL}") +elseif(EXISTS "${THIRDPARTY_DIR}/${PAIMON_LUMINA_PKG_NAME}") + set_urls(LUMINA_SOURCE_URL "${THIRDPARTY_DIR}/${PAIMON_LUMINA_PKG_NAME}") +else() + set_urls(LUMINA_SOURCE_URL + "https://paimon-cpp.oss-cn-beijing.aliyuncs.com/thirdparty/lumina/lumina_release-${PAIMON_LUMINA_BUILD_VERSION}.tar.gz" + ) +endif() + +if(DEFINED ENV{PAIMON_BOOST_URL}) + set(BOOST_SOURCE_URL "$ENV{PAIMON_BOOST_URL}") +elseif(EXISTS "${THIRDPARTY_DIR}/${PAIMON_BOOST_PKG_NAME}") + set_urls(BOOST_SOURCE_URL "${THIRDPARTY_DIR}/${PAIMON_BOOST_PKG_NAME}") +else() + set_urls(BOOST_SOURCE_URL + "https://paimon-cpp.oss-cn-beijing.aliyuncs.com/thirdparty/boost/${PAIMON_BOOST_PKG_NAME}" + ) +endif() + if(APPLE) set(JINDOSDK_C_DYNAMIC_LIB_NAME "jindosdk_c.${PAIMON_JINDOSDK_C_BUILD_VERSION}") set(JINDOSDK_C_DYNAMIC_LIB_FILE "lib${JINDOSDK_C_DYNAMIC_LIB_NAME}.dylib") @@ -938,7 +958,7 @@ macro(build_boost) endif() externalproject_add(boost_ep - URL "${THIRDPARTY_DIR}/boost/${PAIMON_BOOST_PKG_NAME}" + URL ${BOOST_SOURCE_URL} URL_HASH "SHA256=${PAIMON_BOOST_BUILD_SHA256_CHECKSUM}" CONFIGURE_COMMAND ${BOOST_PREFIX}/src/boost_ep/bootstrap.sh --with-libraries=date_time,filesystem,iostreams,regex,system,thread,chrono,atomic @@ -1196,6 +1216,42 @@ macro(build_jindosdk_c) endmacro() +macro(build_lumina) + message(STATUS "Installing Lumina from precompiled package") + + set(LUMINA_PREFIX "${CMAKE_CURRENT_BINARY_DIR}/lumina_ep-install") + set(LUMINA_INCLUDE_DIR "${LUMINA_PREFIX}/include") + set(LUMINA_LIB_DIR "${LUMINA_PREFIX}/lib") + set(LUMINA_DYNAMIC_LIB "${LUMINA_LIB_DIR}/liblumina.so") + + externalproject_add(lumina_ep + URL ${LUMINA_SOURCE_URL} + URL_HASH "SHA256=${PAIMON_LUMINA_BUILD_SHA256_CHECKSUM}" + ${THIRDPARTY_LOG_OPTIONS} + CONFIGURE_COMMAND "" + BUILD_COMMAND "" + INSTALL_COMMAND ${CMAKE_COMMAND} -E copy_directory + /artifacts/cpp/linux-x86_64/install-root/usr/local/include + ${LUMINA_INCLUDE_DIR} + COMMAND ${CMAKE_COMMAND} -E copy_if_different + /artifacts/cpp/linux-x86_64/install-root/usr/local/lib/liblumina.so + ${LUMINA_DYNAMIC_LIB} + BUILD_BYPRODUCTS "${LUMINA_DYNAMIC_LIB}") + + file(MAKE_DIRECTORY "${LUMINA_INCLUDE_DIR}") + file(MAKE_DIRECTORY "${LUMINA_LIB_DIR}") + + add_library(lumina::interface SHARED IMPORTED GLOBAL) + set_target_properties(lumina::interface + PROPERTIES IMPORTED_LOCATION "${LUMINA_DYNAMIC_LIB}" + IMPORTED_NO_SONAME TRUE + INTERFACE_INCLUDE_DIRECTORIES + "${LUMINA_INCLUDE_DIR}") + add_dependencies(lumina::interface lumina_ep) + + install(FILES "${LUMINA_DYNAMIC_LIB}" DESTINATION ${CMAKE_INSTALL_LIBDIR}) +endmacro() + macro(build_jindosdk_nextarch) message(STATUS "Building jindosdk-nextarch from local source") @@ -1817,6 +1873,9 @@ if(PAIMON_ENABLE_JINDO) build_jindosdk_c() build_jindosdk_nextarch() endif() +if(PAIMON_ENABLE_LUMINA) + build_lumina() +endif() if(PAIMON_ENABLE_LUCENE) build_boost() build_lucene() diff --git a/docs/source/build_system.rst b/docs/source/build_system.rst index 61739eb3..6a61b68a 100644 --- a/docs/source/build_system.rst +++ b/docs/source/build_system.rst @@ -91,6 +91,8 @@ Paimon provides a set of built-in optional plugins that you can link to as neede - Index plugins: - ``paimon_file_index_shared`` / ``paimon_file_index_static`` + - ``paimon_lumina_index_shared`` / ``paimon_lumina_index_static`` + - ``paimon_lucene_index_shared`` / ``paimon_lucene_index_static`` .. note:: diff --git a/docs/source/building.rst b/docs/source/building.rst index 68efa5ca..71df728d 100644 --- a/docs/source/building.rst +++ b/docs/source/building.rst @@ -124,6 +124,8 @@ boolean flags to ``cmake``. * ``-DPAIMON_ENABLE_ORC=ON``: Paimon integration with Apache ORC * ``-DPAIMON_ENABLE_AVRO=ON``: Apache Avro libraries and Paimon integration * ``-DPAIMON_ENABLE_JINDO=ON``: Support for Alibaba Jindo filesystems +* ``-DPAIMON_ENABLE_LUMINA=ON``: Support for the Lumina vector index. +* ``-DPAIMON_ENABLE_LUCENE=ON``: Support for Lucene full-text search indexes Third-party dependency source ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/docs/source/user_guide/global_index.rst b/docs/source/user_guide/global_index.rst index 02a13c99..be585329 100644 --- a/docs/source/user_guide/global_index.rst +++ b/docs/source/user_guide/global_index.rst @@ -26,6 +26,7 @@ Paimon supports multiple global index types: - **BTree Index**: An efficient index based on multi-level SST files for scalar column lookups. - **Range Bitmap Index**: A range bitmap index optimized for range predicates on ordered scalar columns. Extends the bitmap approach by encoding value ordering, enabling efficient less-than, greater-than, and range conditions. - **Lucene Index**: A full-text search index powered by Lucene++. Supports tokenized text search with multiple modes including match-all, match-any, phrase, prefix, and wildcard queries. +- **Vector Index (Lumina)**: An approximate nearest neighbor (ANN) index powered by Lumina for vector similarity search with configurable distance metrics. Global indexes work on top of Data Evolution tables. To use global indexes, your table must have: @@ -85,3 +86,10 @@ search modes including match-all, match-any, phrase, prefix, and wildcard querie - **Environment Variable**: ``PAIMON_JIEBA_DICT_DIR`` - **Description**: Specifies the directory containing Jieba dictionary files for Chinese text tokenization. At runtime, the system first checks this environment variable; if not set, it falls back to the compile-time ``JIEBA_TEST_DICT_DIR`` macro (only available in test builds). If neither is available, will fail with an error. + +Vector Index (Lumina) +--------------------- + +An approximate nearest neighbor (ANN) index powered by Lumina for vector similarity search. +It supports high-dimensional vector search with configurable distance metrics and encoding strategies. +For more configurations, refer to the ``docs/reference`` directory in the Lumina release package. diff --git a/src/paimon/global_index/lumina/CMakeLists.txt b/src/paimon/global_index/lumina/CMakeLists.txt new file mode 100644 index 00000000..f85af649 --- /dev/null +++ b/src/paimon/global_index/lumina/CMakeLists.txt @@ -0,0 +1,57 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +if(PAIMON_ENABLE_LUMINA) + set(PAIMON_LUMINA_INDEX lumina_global_index.cpp lumina_global_index_factory.cpp) + + add_paimon_lib(paimon_lumina_index + SOURCES + ${PAIMON_LUMINA_INDEX} + DEPENDENCIES + paimon_shared + lumina::interface + STATIC_LINK_LIBS + arrow + glog + fmt + dl + Threads::Threads + SHARED_LINK_LIBS + lumina::interface + paimon_shared + SHARED_LINK_FLAGS + ${PAIMON_VERSION_SCRIPT_FLAGS}) + + if(PAIMON_BUILD_TESTS) + add_paimon_test(lumina_index_test + SOURCES + lumina_api_test.cpp + lumina_file_io_test.cpp + lumina_global_index_test.cpp + EXTRA_INCLUDES + ${LUMINA_INCLUDE_DIR} + STATIC_LINK_LIBS + paimon_shared + test_utils_static + "-Wl,--whole-archive" + paimon_local_file_system_static + paimon_lumina_index_static + "-Wl,--no-whole-archive" + lumina::interface + ${GTEST_LINK_TOOLCHAIN}) + endif() + +endif() diff --git a/src/paimon/global_index/lumina/lumina_api_test.cpp b/src/paimon/global_index/lumina/lumina_api_test.cpp new file mode 100644 index 00000000..d92315d3 --- /dev/null +++ b/src/paimon/global_index/lumina/lumina_api_test.cpp @@ -0,0 +1,222 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "lumina/api/LuminaBuilder.h" +#include "lumina/api/LuminaSearcher.h" +#include "lumina/core/Constants.h" +#include "lumina/core/Status.h" +#include "lumina/core/Types.h" +#include "lumina/extensions/SearchWithFilterExtension.h" +#include "paimon/fs/local/local_file_system.h" +#include "paimon/global_index/lumina/lumina_file_reader.h" +#include "paimon/global_index/lumina/lumina_file_writer.h" +#include "paimon/global_index/lumina/lumina_memory_pool.h" +#include "paimon/testing/utils/testharness.h" +namespace paimon::lumina::test { +class LuminaInterfaceTest : public ::testing::Test { + public: + void SetUp() override {} + void TearDown() override {} + + void WriteAndFlush(const std::string& index_path, + const std::vector<::lumina::core::vector_id_t>& row_ids) const { + auto fs = std::make_shared(); + std::shared_ptr paimon_pool = GetMemoryPool(); + auto pool = std::make_shared(paimon_pool); + ::lumina::core::MemoryResourceConfig memory_resource(pool.get()); + { // create writer + ::lumina::api::BuilderOptions builder_options; + builder_options.Set(::lumina::core::kIndexType, ::lumina::core::kIndexTypeBruteforce) + .Set(::lumina::core::kDimension, 4) + .Set(::lumina::core::kDistanceMetric, ::lumina::core::kDistanceL2) + .Set(::lumina::core::kEncodingType, ::lumina::core::kEncodingRawf32); + auto builder_result = + ::lumina::api::LuminaBuilder::Create(builder_options, memory_resource); + ASSERT_TRUE(builder_result.IsOk()) << builder_result.GetStatus().Message(); + auto writer = std::move(builder_result).TakeValue(); + // pretrain + ASSERT_TRUE(writer.Pretrain(/*data=*/nullptr, /*n=*/0).IsOk()); + // insert data + auto status = writer.InsertBatch(data_vec_.data(), row_ids.data(), 4); + ASSERT_TRUE(status.IsOk()); + // check memory paimon_pool + ASSERT_GT(paimon_pool->CurrentUsage(), 0); + ASSERT_GT(paimon_pool->MaxMemoryUsage(), 0); + // dump index + ASSERT_OK_AND_ASSIGN(std::shared_ptr out, + fs->Create(index_path, /*overwrite=*/false)); + auto file_writer = std::make_unique(out); + ASSERT_TRUE(writer.Dump(std::move(file_writer), ::lumina::api::IOOptions()).IsOk()); + } + // check memory paimon_pool + ASSERT_EQ(paimon_pool->CurrentUsage(), 0); + ASSERT_GT(paimon_pool->MaxMemoryUsage(), 0); + } + + void Search(const std::string& index_path, int32_t topk, + const std::vector<::lumina::core::vector_id_t>& expected_row_ids, + const std::vector& expected_distances, + const std::function& filter = nullptr) const { + ASSERT_EQ(expected_row_ids.size(), expected_distances.size()); + auto fs = std::make_shared(); + std::shared_ptr paimon_pool = GetMemoryPool(); + auto pool = std::make_shared(paimon_pool); + ::lumina::core::MemoryResourceConfig memory_resource(pool.get()); + { + // create reader + ::lumina::api::SearcherOptions searcher_options; + searcher_options.Set(::lumina::core::kIndexType, ::lumina::core::kIndexTypeBruteforce) + .Set(::lumina::core::kDimension, 4); + auto reader_result = + ::lumina::api::LuminaSearcher::Create(searcher_options, memory_resource); + ASSERT_TRUE(reader_result.IsOk()); + auto reader = std::move(reader_result).TakeValue(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr in, fs->Open(index_path)); + auto file_reader = std::make_unique(in); + ASSERT_TRUE(reader.Open(std::move(file_reader), ::lumina::api::IOOptions()).IsOk()); + // check meta + auto meta = reader.GetMeta(); + ASSERT_EQ(meta.dim, 4); + ASSERT_EQ(meta.count, 4); + + // check memory paimon_pool + ASSERT_GT(paimon_pool->CurrentUsage(), 0); + ASSERT_GT(paimon_pool->MaxMemoryUsage(), 0); + + auto search = [&](int32_t parallel_number) { + ::lumina::api::Query query(query_.data(), query_.size()); + ::lumina::api::SearchOptions search_options; + search_options.Set(::lumina::core::kTopK, topk); + if (parallel_number > 0) { + search_options.Set(::lumina::core::kSearchParallelNumber, parallel_number); + } + + if (!filter) { + auto search_result = reader.Search(query, search_options, *pool); + ASSERT_TRUE(search_result.IsOk()) << search_result.GetStatus().Message(); + CheckResult(search_result.Value().topk, expected_row_ids, expected_distances); + } else { + search_options.Set(::lumina::core::kSearchThreadSafeFilter, true); + ::lumina::extensions::SearchWithFilterExtension reader_with_filter; + ASSERT_TRUE(reader.Attach(reader_with_filter).IsOk()); + auto search_result = + reader_with_filter.SearchWithFilter(query, filter, search_options, *pool); + ASSERT_TRUE(search_result.IsOk()) << search_result.GetStatus().Message(); + CheckResult(search_result.Value().topk, expected_row_ids, expected_distances); + } + + // TODO(xinyu.lxy): check memory paimon_pool, current memory use = query mem + + // reader mem + ASSERT_GT(paimon_pool->CurrentUsage(), 0); + ASSERT_GT(paimon_pool->MaxMemoryUsage(), 0); + }; + + // single thread search result + search(-1); + + // multi thread search result + search(4); + } + // check memory paimon_pool + ASSERT_EQ(paimon_pool->CurrentUsage(), 0); + ASSERT_GT(paimon_pool->MaxMemoryUsage(), 0); + } + + void CheckResult(const std::vector<::lumina::api::LuminaSearcher::SearchHit>& search_result, + const std::vector<::lumina::core::vector_id_t>& expected_row_ids, + const std::vector& expected_distances) const { + ASSERT_EQ(search_result.size(), expected_row_ids.size()); + for (size_t i = 0; i < search_result.size(); i++) { + ASSERT_EQ(expected_row_ids[i], search_result[i].id); + ASSERT_NEAR(expected_distances[i], search_result[i].distance, 0.01); + } + } + + private: + std::vector data_vec_ = { + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 1.0f, + 1.0f, 0.0f, 1.0f, 0.0f, 1.0f, 1.0f, 1.0f, 1.0f, + }; + std::vector query_ = {1.0f, 1.0f, 1.0f, 1.1f}; +}; + +TEST_F(LuminaInterfaceTest, TestSimple) { + auto dir = paimon::test::UniqueTestDirectory::Create("local"); + std::string index_path = dir->Str() + "/lumina_test.index"; + + // write index + std::vector<::lumina::core::vector_id_t> row_ids = {0l, 1l, 2l, 3l}; + WriteAndFlush(index_path, row_ids); + + // read index + std::vector<::lumina::core::vector_id_t> expected_row_ids = {3l, 1l, 2l, 0l}; + std::vector expected_distances = {0.01f, 2.01f, 2.21f, 4.21f}; + Search(index_path, /*topk=*/4, expected_row_ids, expected_distances); +} + +TEST_F(LuminaInterfaceTest, TestWithDocIdGap) { + auto dir = paimon::test::UniqueTestDirectory::Create("local"); + std::string index_path = dir->Str() + "/lumina_test.index"; + + // write index + std::vector<::lumina::core::vector_id_t> row_ids = {0l, 2l, 4l, 6l}; + WriteAndFlush(index_path, row_ids); + + // read index + std::vector<::lumina::core::vector_id_t> expected_row_ids = {6l, 2l, 4l, 0l}; + std::vector expected_distances = {0.01f, 2.01f, 2.21f, 4.21f}; + Search(index_path, /*topk=*/4, expected_row_ids, expected_distances); +} + +TEST_F(LuminaInterfaceTest, TestWithSmallTopk) { + auto dir = paimon::test::UniqueTestDirectory::Create("local"); + std::string index_path = dir->Str() + "/lumina_test.index"; + + // write index + std::vector<::lumina::core::vector_id_t> row_ids = {0l, 1l, 2l, 3l}; + WriteAndFlush(index_path, row_ids); + + // read index + std::vector<::lumina::core::vector_id_t> expected_row_ids = {3l, 1l, 2l}; + std::vector expected_distances = {0.01f, 2.01f, 2.21f}; + Search(index_path, /*topk=*/3, expected_row_ids, expected_distances); +} + +TEST_F(LuminaInterfaceTest, TestWithFilter) { + auto dir = paimon::test::UniqueTestDirectory::Create("local"); + std::string index_path = dir->Str() + "/lumina_test.index"; + + // write index + std::vector<::lumina::core::vector_id_t> row_ids = {0l, 1l, 2l, 3l}; + WriteAndFlush(index_path, row_ids); + + // read index + { + std::vector<::lumina::core::vector_id_t> expected_row_ids = {1l, 2l}; + std::vector expected_distances = {2.01f, 2.21f}; + auto filter = [](::lumina::core::vector_id_t id) -> bool { return id < 3; }; + Search(index_path, /*topk=*/2, expected_row_ids, expected_distances, filter); + } + { + std::vector<::lumina::core::vector_id_t> expected_row_ids = {1l, 2l, 0l}; + std::vector expected_distances = {2.01f, 2.21f, 4.21f}; + auto filter = [](::lumina::core::vector_id_t id) -> bool { return id < 3; }; + Search(index_path, /*topk=*/4, expected_row_ids, expected_distances, filter); + } +} + +} // namespace paimon::lumina::test diff --git a/src/paimon/global_index/lumina/lumina_file_io_test.cpp b/src/paimon/global_index/lumina/lumina_file_io_test.cpp new file mode 100644 index 00000000..571703c6 --- /dev/null +++ b/src/paimon/global_index/lumina/lumina_file_io_test.cpp @@ -0,0 +1,116 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include + +#include "paimon/global_index/lumina/lumina_file_reader.h" +#include "paimon/global_index/lumina/lumina_file_writer.h" +#include "paimon/testing/utils/testharness.h" +namespace paimon::lumina::test { +class LuminaFileIOTest : public ::testing::Test { + public: + void SetUp() override {} + void TearDown() override {} +}; + +TEST_F(LuminaFileIOTest, TestSimple) { + auto check_write_and_read = [](uint64_t max_write_size, uint64_t max_read_size) { + std::string content = "Hello World."; + auto dir = paimon::test::UniqueTestDirectory::Create("local"); + auto fs = dir->GetFileSystem(); + std::string index_path = dir->Str() + "/lumina_test.index"; + // write content + ASSERT_OK_AND_ASSIGN(std::shared_ptr out, + fs->Create(index_path, /*overwrite=*/false)); + auto writer = std::make_shared(out); + writer->max_write_size_ = max_write_size; + ASSERT_EQ(writer->GetLength().Value(), 0); + ASSERT_TRUE(writer->Write(content.data(), content.length()).IsOk()); + ASSERT_TRUE(writer->Close().IsOk()); + + // check file exist + ASSERT_OK_AND_ASSIGN(bool exist, fs->Exists(index_path)); + ASSERT_TRUE(exist); + ASSERT_OK_AND_ASSIGN(std::unique_ptr file_status, + fs->GetFileStatus(index_path)); + ASSERT_FALSE(file_status->IsDir()); + ASSERT_EQ(file_status->GetLen(), content.length()); + + // read content + ASSERT_OK_AND_ASSIGN(std::shared_ptr in, fs->Open(index_path)); + auto reader = std::make_shared(in); + reader->max_read_size_ = max_read_size; + ASSERT_EQ(reader->GetLength().Value(), content.length()); + ASSERT_EQ(reader->GetPosition().Value(), 0); + std::string read_content(content.size(), 0); + ASSERT_TRUE(reader->Read(read_content.data(), read_content.size()).IsOk()); + ASSERT_EQ(read_content, content); + ASSERT_EQ(reader->GetPosition().Value(), content.size()); + + // test seek + ASSERT_TRUE(reader->Seek(2).IsOk()); + std::string read_content2(3, 0); + ASSERT_TRUE(reader->Read(read_content2.data(), read_content2.size()).IsOk()); + ASSERT_EQ(read_content2, "llo"); + ASSERT_EQ(reader->GetPosition().Value(), 5); + ASSERT_TRUE(reader->Close().IsOk()); + }; + + check_write_and_read(LuminaFileWriter::kMaxWriteSize, LuminaFileReader::kMaxReadSize); + check_write_and_read(2, 2); +} + +TEST_F(LuminaFileIOTest, TestReadAsync) { + auto dir = paimon::test::UniqueTestDirectory::Create("local"); + auto fs = dir->GetFileSystem(); + std::string index_path = dir->Str() + "/lumina_test.index"; + std::string content = "abcdefghijk"; + ASSERT_OK(fs->WriteFile(index_path, content, /*overwrite*/ false)); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr in, fs->Open(index_path)); + auto reader = std::make_shared(in); + + auto check_read_result = [&](uint64_t max_read_size, std::string& read_content) { + reader->max_read_size_ = max_read_size; + bool read_finished = false; + std::promise promise; + std::future future = promise.get_future(); + auto callback = [&](::lumina::core::Status status) { + if (status.IsOk()) { + read_finished = true; + promise.set_value(10); + } else { + read_finished = true; + promise.set_value(20); + } + }; + reader->ReadAsync(read_content.data(), read_content.size(), /*offset=*/0, callback); + + ASSERT_EQ(future.get(), 10); + ASSERT_TRUE(read_finished); + ASSERT_EQ(content.substr(0, read_content.size()), read_content); + }; + + for (auto max_read_size : {1, 2, 3, 4, 5, 10, 100}) { + std::string read_content(content.size(), '\0'); + check_read_result(max_read_size, read_content); + // test read empty + std::string empty_content; + check_read_result(max_read_size, empty_content); + } +} +} // namespace paimon::lumina::test diff --git a/src/paimon/global_index/lumina/lumina_file_reader.h b/src/paimon/global_index/lumina/lumina_file_reader.h new file mode 100644 index 00000000..db6d2fbb --- /dev/null +++ b/src/paimon/global_index/lumina/lumina_file_reader.h @@ -0,0 +1,141 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include + +#include "fmt/format.h" +#include "lumina/io/FileReader.h" +#include "paimon/fs/file_system.h" +#include "paimon/global_index/lumina/lumina_utils.h" +namespace paimon::lumina { +class LuminaFileReader : public ::lumina::io::FileReader { + public: + explicit LuminaFileReader(const std::shared_ptr& in) : in_(in) {} + ~LuminaFileReader() override = default; + + ::lumina::core::Result GetLength() const noexcept override { + Result length_result = in_->Length(); + if (!length_result.ok()) { + return ::lumina::core::Result::Err( + PaimonToLuminaStatus(length_result.status())); + } + return ::lumina::core::Result::Ok(length_result.value()); + } + + ::lumina::core::Result GetPosition() const noexcept override { + Result pos_result = in_->GetPos(); + if (!pos_result.ok()) { + return ::lumina::core::Result::Err(PaimonToLuminaStatus(pos_result.status())); + } + return ::lumina::core::Result::Ok(static_cast(pos_result.value())); + } + + ::lumina::core::Status Seek(uint64_t position) noexcept override { + return PaimonToLuminaStatus(in_->Seek(position, SeekOrigin::FS_SEEK_SET)); + } + + ::lumina::core::Status Read(char* data, uint64_t size) noexcept override { + uint64_t total_read_size = 0; + while (total_read_size < size) { + uint64_t current_read_size = std::min(size - total_read_size, max_read_size_); + Result read_result = + in_->Read(data + total_read_size, static_cast(current_read_size)); + if (!read_result.ok()) { + return PaimonToLuminaStatus(read_result.status()); + } + if (static_cast(read_result.value()) != current_read_size) { + return ::lumina::core::Status( + ::lumina::core::ErrorCode::IoError, + fmt::format("expect read len {} mismatch actual read len {}", current_read_size, + read_result.value())); + } + total_read_size += current_read_size; + } + return ::lumina::core::Status::Ok(); + } + + void ReadAsync(char* data, uint64_t size, uint64_t offset, + std::function call_back) noexcept override { + if (size == 0) { + call_back(::lumina::core::Status::Ok()); + return; + } + + struct ReadContext { + char* current_data; + uint64_t remaining; + uint64_t current_offset; + std::function final_call_back; + std::shared_ptr in; + }; + + auto ctx = std::make_shared( + ReadContext{data, size, offset, std::move(call_back), in_}); + + // recursive lambda to read next chunk + std::function read_next; + read_next = [ctx, max_read_size = max_read_size_, &read_next]() { + if (ctx->remaining == 0) { + // all done + ctx->final_call_back(::lumina::core::Status::Ok()); + return; + } + + // determine this chunk's size + uint64_t chunk_size = std::min(ctx->remaining, max_read_size); + auto safe_size = static_cast(chunk_size); + + // issue async read for this chunk + ctx->in->ReadAsync(ctx->current_data, safe_size, ctx->current_offset, + [ctx, safe_size, read_next](const Status& status) { + if (!status.ok()) { + // propagate error immediately + ctx->final_call_back(PaimonToLuminaStatus(status)); + return; + } + // advance pointers and counters + ctx->current_data += safe_size; + ctx->current_offset += safe_size; + ctx->remaining -= safe_size; + + // continue with next chunk + read_next(); + }); + }; + + // start the first read + read_next(); + } + + ::lumina::core::Status Close() noexcept override { + return PaimonToLuminaStatus(in_->Close()); + } + + private: + static constexpr uint64_t kMaxReadSize = std::numeric_limits::max(); + + private: + uint64_t max_read_size_ = kMaxReadSize; + std::shared_ptr in_; +}; +} // namespace paimon::lumina diff --git a/src/paimon/global_index/lumina/lumina_file_writer.h b/src/paimon/global_index/lumina/lumina_file_writer.h new file mode 100644 index 00000000..16323778 --- /dev/null +++ b/src/paimon/global_index/lumina/lumina_file_writer.h @@ -0,0 +1,78 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include + +#include "fmt/format.h" +#include "lumina/io/FileWriter.h" +#include "paimon/fs/file_system.h" +#include "paimon/global_index/lumina/lumina_utils.h" +namespace paimon::lumina { +class LuminaFileWriter : public ::lumina::io::FileWriter { + public: + explicit LuminaFileWriter(const std::shared_ptr& out) : out_(out) {} + ~LuminaFileWriter() override = default; + + ::lumina::core::Result GetLength() const noexcept override { + Result pos_result = out_->GetPos(); + if (!pos_result.ok()) { + return ::lumina::core::Result::Err(PaimonToLuminaStatus(pos_result.status())); + } + return ::lumina::core::Result::Ok(static_cast(pos_result.value())); + } + + ::lumina::core::Status Write(const char* data, uint64_t size) noexcept override { + uint64_t total_write_size = 0; + while (total_write_size < size) { + uint64_t current_write_size = std::min(size - total_write_size, max_write_size_); + Result write_result = + out_->Write(data + total_write_size, static_cast(current_write_size)); + if (!write_result.ok()) { + return PaimonToLuminaStatus(write_result.status()); + } + if (static_cast(write_result.value()) != current_write_size) { + return ::lumina::core::Status( + ::lumina::core::ErrorCode::IoError, + fmt::format("expect write len {} mismatch actual write len {}", + current_write_size, write_result.value())); + } + total_write_size += current_write_size; + } + return ::lumina::core::Status::Ok(); + } + + ::lumina::core::Status Close() noexcept override { + auto status = out_->Flush(); + if (!status.ok()) { + return PaimonToLuminaStatus(status); + } + return PaimonToLuminaStatus(out_->Close()); + } + + private: + static constexpr uint64_t kMaxWriteSize = std::numeric_limits::max(); + + private: + uint64_t max_write_size_ = kMaxWriteSize; + std::shared_ptr out_; +}; +} // namespace paimon::lumina diff --git a/src/paimon/global_index/lumina/lumina_global_index.cpp b/src/paimon/global_index/lumina/lumina_global_index.cpp new file mode 100644 index 00000000..5976acd0 --- /dev/null +++ b/src/paimon/global_index/lumina/lumina_global_index.cpp @@ -0,0 +1,376 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/global_index/lumina/lumina_global_index.h" + +#include + +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "lumina/api/Dataset.h" +#include "lumina/api/LuminaBuilder.h" +#include "lumina/api/LuminaSearcher.h" +#include "lumina/api/OptionsNormalize.h" +#include "lumina/core/Constants.h" +#include "lumina/core/Status.h" +#include "lumina/core/Types.h" +#include "paimon/common/global_index/global_index_utils.h" +#include "paimon/common/utils/options_utils.h" +#include "paimon/common/utils/rapidjson_util.h" +#include "paimon/common/utils/string_utils.h" +#include "paimon/global_index/bitmap_scored_global_index_result.h" +#include "paimon/global_index/lumina/lumina_file_reader.h" +#include "paimon/global_index/lumina/lumina_file_writer.h" +#include "paimon/global_index/lumina/lumina_utils.h" +namespace paimon::lumina { +#define CHECK_NOT_NULL(pointer, error_msg) \ + do { \ + if (!(pointer)) { \ + return Status::Invalid(error_msg); \ + } \ + } while (0) + +Result> LuminaGlobalIndex::CreateWriter( + const std::string& field_name, ::ArrowSchema* arrow_schema, + const std::shared_ptr& file_writer, + const std::shared_ptr& pool) const { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr arrow_type, + arrow::ImportType(arrow_schema)); + // check data type + auto struct_type = std::dynamic_pointer_cast(arrow_type); + CHECK_NOT_NULL(struct_type, "arrow schema must be struct type when create LuminaIndexWriter"); + auto index_field = struct_type->GetFieldByName(field_name); + CHECK_NOT_NULL(index_field, + fmt::format("field {} not exist in arrow schema when create LuminaIndexWriter", + field_name)); + auto list_type = std::dynamic_pointer_cast(index_field->type()); + CHECK_NOT_NULL(list_type, "field type must be list[float] when create LuminaIndexWriter"); + if (list_type->value_type()->id() != arrow::Type::type::FLOAT) { + return Status::Invalid("field type must be list[float] when create LuminaIndexWriter"); + } + + // check options + auto lumina_options = + OptionsUtils::FetchOptionsWithPrefix(LuminaDefines::kOptionKeyPrefix, options_); + PAIMON_ASSIGN_OR_RAISE(uint32_t dimension, + OptionsUtils::GetValueFromMap( + lumina_options, std::string(::lumina::core::kDimension))); + + PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( + ::lumina::api::BuilderOptions builder_options, + ::lumina::api::NormalizeBuilderOptions(std::unordered_map( + lumina_options.begin(), lumina_options.end()))); + auto lumina_pool = std::make_shared(pool); + return std::make_shared( + field_name, arrow_type, dimension, file_writer, std::move(builder_options), + ::lumina::api::IOOptions(), lumina_options, lumina_pool); +} + +Result LuminaIndexReader::GetIndexInfo( + const GlobalIndexIOMeta& io_meta) { + auto meta_bytes = io_meta.metadata; + if (!meta_bytes) { + return Status::Invalid("Lumina global index must have meta data"); + } + std::map lumina_write_options; + PAIMON_RETURN_NOT_OK(RapidJsonUtil::FromJsonString( + std::string(meta_bytes->data(), meta_bytes->size()), &lumina_write_options)); + + // check options + PAIMON_ASSIGN_OR_RAISE(uint32_t dimension, + OptionsUtils::GetValueFromMap( + lumina_write_options, std::string(::lumina::core::kDimension))); + PAIMON_ASSIGN_OR_RAISE(std::string index_type, + OptionsUtils::GetValueFromMap( + lumina_write_options, std::string(::lumina::core::kIndexType))); + PAIMON_ASSIGN_OR_RAISE(std::string distance_type_str, + OptionsUtils::GetValueFromMap( + lumina_write_options, std::string(::lumina::core::kDistanceMetric))); + VectorSearch::DistanceType distance_type = VectorSearch::DistanceType::UNKNOWN; + if (distance_type_str == ::lumina::core::kDistanceL2) { + distance_type = VectorSearch::DistanceType::EUCLIDEAN; + } else if (distance_type_str == ::lumina::core::kDistanceCosine) { + distance_type = VectorSearch::DistanceType::COSINE; + } else if (distance_type_str == ::lumina::core::kDistanceInnerProduct) { + distance_type = VectorSearch::DistanceType::INNER_PRODUCT; + } + if (distance_type == VectorSearch::DistanceType::UNKNOWN) { + return Status::Invalid( + fmt::format("invalid distance type {} for lumina", distance_type_str)); + } + return LuminaIndexReader::IndexInfo({dimension, index_type, distance_type}); +} + +Result> LuminaGlobalIndex::CreateReader( + ::ArrowSchema* c_arrow_schema, const std::shared_ptr& file_manager, + const std::vector& files, const std::shared_ptr& pool) const { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr arrow_schema, + arrow::ImportSchema(c_arrow_schema)); + if (files.size() != 1) { + return Status::Invalid("lumina index only has one index file per shard"); + } + const auto& io_meta = files[0]; + // check data type + if (arrow_schema->num_fields() != 1) { + return Status::Invalid("LuminaGlobalIndex now only support one field"); + } + auto index_field = arrow_schema->field(0); + auto list_type = std::dynamic_pointer_cast(index_field->type()); + CHECK_NOT_NULL(list_type, "field type must be list[float] when create LuminaIndexReader"); + if (list_type->value_type()->id() != arrow::Type::type::FLOAT) { + return Status::Invalid("field type must be list[float] when create LuminaIndexReader"); + } + + // get index info from meta + PAIMON_ASSIGN_OR_RAISE(LuminaIndexReader::IndexInfo index_info, + LuminaIndexReader::GetIndexInfo(io_meta)); + + auto lumina_pool = std::make_shared(pool); + ::lumina::core::MemoryResourceConfig memory_resource(lumina_pool.get()); + + auto lumina_options = + OptionsUtils::FetchOptionsWithPrefix(LuminaDefines::kOptionKeyPrefix, options_); + lumina_options[std::string(::lumina::core::kDimension)] = std::to_string(index_info.dimension); + lumina_options[std::string(::lumina::core::kIndexType)] = index_info.index_type; + + PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( + ::lumina::api::SearcherOptions searcher_options, + ::lumina::api::NormalizeSearcherOptions(std::unordered_map( + lumina_options.begin(), lumina_options.end()))); + + PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( + ::lumina::api::LuminaSearcher lumina_searcher, + ::lumina::api::LuminaSearcher::Create(searcher_options, memory_resource)); + auto searcher = std::make_unique<::lumina::api::LuminaSearcher>(std::move(lumina_searcher)); + // get input stream and open index + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr in, + file_manager->GetInputStream(io_meta.file_path)); + auto lumina_file_reader = std::make_unique(in); + PAIMON_RETURN_NOT_OK_FROM_LUMINA( + searcher->Open(std::move(lumina_file_reader), ::lumina::api::IOOptions())); + + // check meta + if (searcher->GetMeta().dim != index_info.dimension) { + return Status::Invalid( + fmt::format("lumina index dimension {} mismatch dimension {} in io meta", + searcher->GetMeta().dim, index_info.dimension)); + } + auto searcher_with_filter = std::make_unique<::lumina::extensions::SearchWithFilterExtension>(); + PAIMON_RETURN_NOT_OK_FROM_LUMINA(searcher->Attach(*searcher_with_filter)); + return std::make_shared(index_info, std::move(searcher), + std::move(searcher_with_filter), lumina_pool); +} + +class LuminaDataset : public ::lumina::api::Dataset { + public: + LuminaDataset(int64_t element_count, uint32_t dimension, + const std::vector>& array_vec) + : element_count_(element_count), dimension_(dimension), array_vec_(array_vec) {} + + uint32_t Dim() const noexcept override { + return dimension_; + } + uint64_t TotalSize() const noexcept override { + return element_count_; + } + + ::lumina::core::Result GetNextBatch( + std::vector& vector_buffer, + std::vector<::lumina::core::vector_id_t>& id_buffer) noexcept override { + if (cursor_ >= array_vec_.size()) { + return ::lumina::core::Result::Ok(0); + } + auto& value_array = array_vec_[cursor_]; + int64_t value_array_length = value_array->length(); + int64_t element_count = value_array_length / dimension_; + const float* value_ptr = value_array->raw_values(); + vector_buffer.resize(value_array_length); + memcpy(vector_buffer.data(), value_ptr, sizeof(float) * value_array_length); + id_buffer.resize(element_count); + std::iota(id_buffer.begin(), id_buffer.end(), id_); + id_ += element_count; + + // release the array when copy to vector_buffer + value_array.reset(); + cursor_++; + return ::lumina::core::Result::Ok(static_cast(element_count)); + } + + private: + int64_t element_count_; + uint32_t dimension_; + std::vector> array_vec_; + size_t cursor_ = 0; + ::lumina::core::vector_id_t id_ = 0; +}; + +LuminaIndexWriter::LuminaIndexWriter(const std::string& field_name, + const std::shared_ptr& arrow_type, + uint32_t dimension, + const std::shared_ptr& file_manager, + ::lumina::api::BuilderOptions&& builder_options, + ::lumina::api::IOOptions&& io_options, + const std::map& lumina_options, + const std::shared_ptr& pool) + : pool_(pool), + field_name_(field_name), + arrow_type_(arrow_type), + dimension_(dimension), + file_manager_(file_manager), + builder_options_(std::move(builder_options)), + io_options_(std::move(io_options)), + lumina_options_(lumina_options) {} + +Status LuminaIndexWriter::AddBatch(::ArrowArray* arrow_array, + std::vector&& relative_row_ids) { + PAIMON_RETURN_NOT_OK( + GlobalIndexUtils::CheckRelativeRowIds(arrow_array, relative_row_ids, count_)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, + arrow::ImportArray(arrow_array, arrow_type_)); + if (array->null_count() != 0) { + return Status::Invalid("arrow_array in LuminaIndexWriter is invalid, must not null"); + } + auto struct_array = std::dynamic_pointer_cast(array); + CHECK_NOT_NULL(struct_array, "invalid input array in LuminaIndexWriter, must be struct array"); + auto field_array = struct_array->GetFieldByName(field_name_); + CHECK_NOT_NULL( + field_array, + fmt::format("invalid input array in LuminaIndexWriter, field {} not in input array", + field_name_)); + int64_t field_length = field_array->length(); + auto list_field_array = std::dynamic_pointer_cast(field_array); + CHECK_NOT_NULL(list_field_array, + "invalid input array in LuminaIndexWriter, field array must be list array"); + auto value_array = std::dynamic_pointer_cast(list_field_array->values()); + CHECK_NOT_NULL( + value_array, + "invalid input array in LuminaIndexWriter, field value array must be float array"); + if (value_array->null_count() != 0) { + return Status::Invalid("field value array in LuminaIndexWriter is invalid, must not null"); + } + if (value_array->length() != field_length * dimension_) { + return Status::Invalid(fmt::format( + "invalid input array in LuminaIndexWriter, length of field array [{}] multiplied " + "dimension [{}] must match length of field value array [{}]", + field_length, dimension_, value_array->length())); + } + count_ += array->length(); + array_vec_.push_back(std::move(value_array)); + return Status::OK(); +} + +Result> LuminaIndexWriter::Finish() { + ::lumina::core::MemoryResourceConfig memory_resource(pool_.get()); + PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( + ::lumina::api::LuminaBuilder builder, + ::lumina::api::LuminaBuilder::Create(builder_options_, memory_resource)); + // pretrain + LuminaDataset dataset1(count_, dimension_, array_vec_); + PAIMON_RETURN_NOT_OK_FROM_LUMINA(builder.PretrainFrom(dataset1)); + + // insert data + LuminaDataset dataset2(count_, dimension_, array_vec_); + std::vector>().swap(array_vec_); + PAIMON_RETURN_NOT_OK_FROM_LUMINA(builder.InsertFrom(dataset2)); + + // dump index + PAIMON_ASSIGN_OR_RAISE(std::string index_file_name, + file_manager_->NewFileName(LuminaDefines::kIdentifier)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr out, + file_manager_->NewOutputStream(index_file_name)) + auto file_writer = std::make_unique(out); + PAIMON_RETURN_NOT_OK_FROM_LUMINA(builder.Dump(std::move(file_writer), io_options_)); + // prepare GlobalIndexIOMeta + PAIMON_ASSIGN_OR_RAISE(int64_t file_size, file_manager_->GetFileSize(index_file_name)); + std::string options_json; + PAIMON_RETURN_NOT_OK(RapidJsonUtil::ToJsonString(lumina_options_, &options_json)); + auto meta_bytes = std::make_shared(options_json, pool_->GetPaimonPool().get()); + GlobalIndexIOMeta meta(file_manager_->ToPath(index_file_name), file_size, + /*metadata=*/meta_bytes); + return std::vector({meta}); +} + +LuminaIndexReader::LuminaIndexReader( + const LuminaIndexReader::IndexInfo& index_info, + std::unique_ptr<::lumina::api::LuminaSearcher>&& searcher, + std::unique_ptr<::lumina::extensions::SearchWithFilterExtension>&& searcher_with_filter, + const std::shared_ptr& pool) + : index_info_(index_info), + pool_(pool), + searcher_(std::move(searcher)), + searcher_with_filter_(std::move(searcher_with_filter)) {} + +Result> LuminaIndexReader::VisitVectorSearch( + const std::shared_ptr& vector_search) { + if (vector_search->predicate) { + return Status::NotImplemented("lumina index not support predicate in VisitVectorSearch"); + } + if (vector_search->distance_type && + vector_search->distance_type.value() != index_info_.distance_type) { + return Status::Invalid("distance type for index and search not match"); + } + if (vector_search->query.size() != index_info_.dimension) { + return Status::Invalid("dimension for index and search not match"); + } + + auto lumina_options = OptionsUtils::FetchOptionsWithPrefix(LuminaDefines::kOptionKeyPrefix, + vector_search->options); + auto index_type_iter = lumina_options.find(std::string(::lumina::core::kIndexType)); + if (index_type_iter != lumina_options.end() && + index_type_iter->second != index_info_.index_type) { + return Status::Invalid("index type for index and search not match"); + } + + lumina_options[std::string(::lumina::core::kTopK)] = std::to_string(vector_search->limit); + lumina_options[std::string(::lumina::core::kSearchThreadSafeFilter)] = "true"; + PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( + ::lumina::api::SearchOptions search_options, + ::lumina::api::NormalizeSearchOptions(index_info_.index_type, + std::unordered_map( + lumina_options.begin(), lumina_options.end()))); + + ::lumina::api::Query lumina_query(vector_search->query.data(), vector_search->query.size()); + ::lumina::api::LuminaSearcher::SearchResult search_result; + if (!vector_search->pre_filter) { + PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA(search_result, + searcher_->Search(lumina_query, search_options, *pool_)); + } else { + auto lumina_filter = [filter = vector_search->pre_filter]( + ::lumina::core::vector_id_t id) -> bool { return filter(id); }; + PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( + search_result, searcher_with_filter_->SearchWithFilter(lumina_query, lumina_filter, + search_options, *pool_)); + } + + // prepare BitmapScoredGlobalIndexResult + std::map id_to_score; + for (const auto& [id, score] : search_result.topk) { + id_to_score[id] = score; + } + + RoaringBitmap64 bitmap; + std::vector scores; + scores.reserve(id_to_score.size()); + for (const auto& [id, score] : id_to_score) { + bitmap.Add(id); + scores.push_back(score); + } + return std::make_shared(std::move(bitmap), std::move(scores)); +} + +} // namespace paimon::lumina diff --git a/src/paimon/global_index/lumina/lumina_global_index.h b/src/paimon/global_index/lumina/lumina_global_index.h new file mode 100644 index 00000000..281f5185 --- /dev/null +++ b/src/paimon/global_index/lumina/lumina_global_index.h @@ -0,0 +1,209 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "lumina/api/LuminaSearcher.h" +#include "lumina/api/Options.h" +#include "lumina/extensions/SearchWithFilterExtension.h" +#include "paimon/global_index/bitmap_global_index_result.h" +#include "paimon/global_index/global_indexer.h" +#include "paimon/global_index/lumina/lumina_memory_pool.h" +#include "paimon/global_index/lumina/lumina_utils.h" + +namespace paimon::lumina { +/// @note When enabling the lumina global index in `paimon-cpp`, all configuration parameters +/// specific to Lumina **must be prefixed with `lumina.`**. +/// See `docs/reference/OptionsReference.md` in the Lumina release package for more options. +/// +/// Example configurations: +/// +/// - **Index Writer (build-time options):** +/// lumina.index.dimension:1024 +/// lumina.index.type:diskann +/// lumina.distance.metric:inner_product +/// lumina.encoding.type:pq +/// lumina.encoding.pq.m:256 +/// lumina.diskann.build.thread_count:64 +/// lumina.diskann.build.ef_construction:1024 +/// lumina.diskann.build.neighbor_count:64 +/// +/// - **Index Reader:** +/// No configuration required at load time — settings are stored in the index metadata, +/// and the this plugin will automatically infer and apply them during loading. +/// +/// - **Vector Search (query-time options):** +/// lumina.search.parallel_number:5 +/// lumina.diskann.search.beam_width:4 +/// lumina.diskann.search.list_size:1024 +class LuminaGlobalIndex : public GlobalIndexer { + public: + explicit LuminaGlobalIndex(const std::map& options) + : options_(options) {} + + Result> CreateWriter( + const std::string& field_name, ::ArrowSchema* arrow_schema, + const std::shared_ptr& file_writer, + const std::shared_ptr& pool) const override; + + Result> CreateReader( + ::ArrowSchema* arrow_schema, const std::shared_ptr& file_reader, + const std::vector& files, + const std::shared_ptr& pool) const override; + + private: + std::map options_; +}; + +class LuminaIndexWriter : public GlobalIndexWriter { + public: + LuminaIndexWriter(const std::string& field_name, + const std::shared_ptr& arrow_type, uint32_t dimension, + const std::shared_ptr& file_manager, + ::lumina::api::BuilderOptions&& builder_options, + ::lumina::api::IOOptions&& io_options, + const std::map& lumina_options, + const std::shared_ptr& pool); + + Status AddBatch(::ArrowArray* arrow_array, std::vector&& relative_row_ids) override; + + Result> Finish() override; + + private: + int64_t count_ = 0; + std::shared_ptr pool_; + std::string field_name_; + std::shared_ptr arrow_type_; + uint32_t dimension_; + std::shared_ptr file_manager_; + ::lumina::api::BuilderOptions builder_options_; + ::lumina::api::IOOptions io_options_; + std::map lumina_options_; + std::vector> array_vec_; +}; + +class LuminaIndexReader : public GlobalIndexReader { + public: + struct IndexInfo { + uint32_t dimension; + std::string index_type; + VectorSearch::DistanceType distance_type; + }; + + LuminaIndexReader( + const IndexInfo& index_info, std::unique_ptr<::lumina::api::LuminaSearcher>&& searcher, + std::unique_ptr<::lumina::extensions::SearchWithFilterExtension>&& searcher_with_filter, + const std::shared_ptr& pool); + + ~LuminaIndexReader() override { + [[maybe_unused]] auto status = searcher_->Close(); + } + + /// @note `VisitVectorSearch` is thread-safe (not coroutine-safe) while other `VisitXXX` is not + /// thread-safe. + Result> VisitVectorSearch( + const std::shared_ptr& vector_search) override; + + Result> VisitFullTextSearch( + const std::shared_ptr& full_text_search) override { + return std::shared_ptr(); + } + + Result> VisitIsNotNull() override { + return std::shared_ptr(); + } + + Result> VisitIsNull() override { + return std::shared_ptr(); + } + + Result> VisitEqual(const Literal& literal) override { + return std::shared_ptr(); + } + + Result> VisitNotEqual(const Literal& literal) override { + return std::shared_ptr(); + } + + Result> VisitLessThan(const Literal& literal) override { + return std::shared_ptr(); + } + + Result> VisitLessOrEqual(const Literal& literal) override { + return std::shared_ptr(); + } + + Result> VisitGreaterThan(const Literal& literal) override { + return std::shared_ptr(); + } + + Result> VisitGreaterOrEqual( + const Literal& literal) override { + return std::shared_ptr(); + } + + Result> VisitIn( + const std::vector& literals) override { + return std::shared_ptr(); + } + + Result> VisitNotIn( + const std::vector& literals) override { + return std::shared_ptr(); + } + + Result> VisitStartsWith(const Literal& prefix) override { + return std::shared_ptr(); + } + + Result> VisitEndsWith(const Literal& suffix) override { + return std::shared_ptr(); + } + + Result> VisitContains(const Literal& literal) override { + return std::shared_ptr(); + } + + Result> VisitLike(const Literal& literal) override { + return std::shared_ptr(); + } + + bool IsThreadSafe() const override { + return true; + } + + std::string GetIndexType() const override { + return LuminaDefines::kIdentifier; + } + + static Result GetIndexInfo(const GlobalIndexIOMeta& io_meta); + + private: + LuminaIndexReader::IndexInfo index_info_; + std::shared_ptr pool_; + std::unique_ptr<::lumina::api::LuminaSearcher> searcher_; + std::unique_ptr<::lumina::extensions::SearchWithFilterExtension> searcher_with_filter_; +}; +} // namespace paimon::lumina diff --git a/src/paimon/global_index/lumina/lumina_global_index_factory.cpp b/src/paimon/global_index/lumina/lumina_global_index_factory.cpp new file mode 100644 index 00000000..e27837da --- /dev/null +++ b/src/paimon/global_index/lumina/lumina_global_index_factory.cpp @@ -0,0 +1,38 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/global_index/lumina/lumina_global_index_factory.h" + +#include +#include +#include +#include + +#include "paimon/global_index/lumina/lumina_global_index.h" +namespace paimon::lumina { + +const char LuminaGlobalIndexFactory::IDENTIFIER[] = "lumina-global"; + +Result> LuminaGlobalIndexFactory::Create( + const std::map& options) const { + return std::make_unique(options); +} + +REGISTER_PAIMON_FACTORY(LuminaGlobalIndexFactory); + +} // namespace paimon::lumina diff --git a/src/paimon/global_index/lumina/lumina_global_index_factory.h b/src/paimon/global_index/lumina/lumina_global_index_factory.h new file mode 100644 index 00000000..582cd001 --- /dev/null +++ b/src/paimon/global_index/lumina/lumina_global_index_factory.h @@ -0,0 +1,41 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include + +#include "paimon/global_index/global_indexer_factory.h" + +namespace paimon::lumina { +/// Factory for creating lumina global indexers. +class LuminaGlobalIndexFactory : public GlobalIndexerFactory { + public: + static const char IDENTIFIER[]; + + const char* Identifier() const override { + return IDENTIFIER; + } + + Result> Create( + const std::map& options) const override; +}; + +} // namespace paimon::lumina diff --git a/src/paimon/global_index/lumina/lumina_global_index_test.cpp b/src/paimon/global_index/lumina/lumina_global_index_test.cpp new file mode 100644 index 00000000..d0bda5fc --- /dev/null +++ b/src/paimon/global_index/lumina/lumina_global_index_test.cpp @@ -0,0 +1,475 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "paimon/global_index/lumina/lumina_global_index.h" + +#include +#include + +#include "arrow/c/bridge.h" +#include "arrow/ipc/api.h" +#include "gtest/gtest.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/date_time_utils.h" +#include "paimon/common/utils/path_util.h" +#include "paimon/common/utils/string_utils.h" +#include "paimon/core/global_index/global_index_file_manager.h" +#include "paimon/core/index/index_path_factory.h" +#include "paimon/fs/local/local_file_system.h" +#include "paimon/global_index/bitmap_scored_global_index_result.h" +#include "paimon/global_index/global_index_result.h" +#include "paimon/predicate/predicate_builder.h" +#include "paimon/testing/utils/testharness.h" +namespace paimon::lumina::test { +class LuminaGlobalIndexTest : public ::testing::Test { + public: + void SetUp() override {} + void TearDown() override {} + + class FakeIndexPathFactory : public IndexPathFactory { + public: + explicit FakeIndexPathFactory(const std::string& index_path) : index_path_(index_path) {} + std::string NewPath() const override { + assert(false); + return ""; + } + std::string ToPath(const std::shared_ptr& file) const override { + assert(false); + return ""; + } + std::string ToPath(const std::string& file_name) const override { + return PathUtil::JoinPath(index_path_, file_name); + } + bool IsExternalPath() const override { + return false; + } + + private: + std::string index_path_; + }; + + std::unique_ptr<::ArrowSchema> CreateArrowSchema( + const std::shared_ptr& data_type) const { + auto c_schema = std::make_unique<::ArrowSchema>(); + EXPECT_TRUE(arrow::ExportType(*data_type, c_schema.get()).ok()); + return c_schema; + } + + Result WriteGlobalIndex(const std::string& index_root, + const std::shared_ptr& data_type, + const std::map& options, + const std::shared_ptr& array, + const Range& expected_range) const { + auto global_index = std::make_shared(options); + auto path_factory = std::make_shared(index_root); + auto file_writer = std::make_shared(fs_, path_factory); + + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr global_writer, + global_index->CreateWriter("f0", CreateArrowSchema(data_type).get(), + file_writer, pool_)); + + ArrowArray c_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); + std::vector row_ids(array->length(), 0); + std::iota(row_ids.begin(), row_ids.end(), 0); + PAIMON_RETURN_NOT_OK(global_writer->AddBatch(&c_array, std::move(row_ids))); + PAIMON_ASSIGN_OR_RAISE(auto result_metas, global_writer->Finish()); + // check meta + EXPECT_EQ(result_metas.size(), 1); + auto file_name = PathUtil::GetName(result_metas[0].file_path); + EXPECT_TRUE(StringUtils::StartsWith(file_name, "lumina-global-index-")); + EXPECT_TRUE(StringUtils::EndsWith(file_name, ".index")); + EXPECT_TRUE(result_metas[0].metadata); + return result_metas[0]; + } + + void CheckResult(const std::shared_ptr& result, + const std::vector& expected_ids, + const std::vector& expected_scores) const { + auto typed_result = std::dynamic_pointer_cast(result); + ASSERT_TRUE(typed_result); + ASSERT_OK_AND_ASSIGN(const RoaringBitmap64* bitmap, typed_result->GetBitmap()); + ASSERT_TRUE(bitmap); + ASSERT_EQ(*(typed_result->GetBitmap().value()), RoaringBitmap64::From(expected_ids)) + << "result=" << (typed_result->GetBitmap().value())->ToString() + << ", expected=" << RoaringBitmap64::From(expected_ids).ToString(); + ASSERT_EQ(typed_result->scores_.size(), expected_scores.size()); + + std::map id_to_score; + for (size_t i = 0; i < expected_ids.size(); i++) { + id_to_score[expected_ids[i]] = expected_scores[i]; + } + std::vector expected_scores_ordered_by_id; + for (const auto& [id, score] : id_to_score) { + expected_scores_ordered_by_id.push_back(score); + } + for (size_t i = 0; i < expected_scores.size(); i++) { + ASSERT_NEAR(typed_result->scores_[i], expected_scores_ordered_by_id[i], 0.01); + } + } + + Result> CreateGlobalIndexReader( + const std::string& index_root, const std::shared_ptr& data_type, + const std::map& options, const GlobalIndexIOMeta& meta) const { + auto global_index = std::make_shared(options); + auto path_factory = std::make_shared(index_root); + auto file_reader = std::make_shared(fs_, path_factory); + return global_index->CreateReader(CreateArrowSchema(data_type).get(), file_reader, {meta}, + pool_); + } + + std::shared_ptr CreateRandomVector(int32_t element_size, + int32_t dimension) const { + int64_t total_values = element_size * dimension; + auto float_builder = std::make_shared(); + arrow::ListBuilder list_builder(arrow::default_memory_pool(), float_builder); + + std::random_device rd; + std::mt19937 gen(rd()); + std::uniform_real_distribution dis(0.0f, 2.0f); + + EXPECT_TRUE(float_builder->Reserve(total_values).ok()); + EXPECT_TRUE(list_builder.Reserve(element_size).ok()); + + for (int64_t i = 0; i < element_size; ++i) { + EXPECT_TRUE(list_builder.Append().ok()); + for (int64_t j = 0; j < dimension; ++j) { + float val = dis(gen); + EXPECT_TRUE(float_builder->Append(val).ok()); + } + } + + std::shared_ptr list_array; + EXPECT_TRUE(list_builder.Finish(&list_array).ok()); + + auto struct_array = arrow::StructArray::Make({list_array}, {"f0"}).ValueOrDie(); + return struct_array; + } + + private: + std::shared_ptr pool_ = GetDefaultPool(); + std::shared_ptr fs_ = std::make_shared(); + std::map options_ = {{"lumina.index.dimension", "4"}, + {"lumina.index.type", "bruteforce"}, + {"lumina.distance.metric", "l2"}, + {"lumina.encoding.type", "rawf32"}, + {"lumina.search.parallel_number", "10"}}; + std::shared_ptr data_type_ = + arrow::struct_({arrow::field("f0", arrow::list(arrow::float32()))}); + std::shared_ptr array_ = arrow::ipc::internal::json::ArrayFromJSON(data_type_, + R"([ + [[0.0, 0.0, 0.0, 0.0]], + [[0.0, 1.0, 0.0, 1.0]], + [[1.0, 0.0, 1.0, 0.0]], + [[1.0, 1.0, 1.0, 1.0]] + ])") + .ValueOrDie(); + std::vector query_ = {1.0f, 1.0f, 1.0f, 1.1f}; +}; + +TEST_F(LuminaGlobalIndexTest, TestSimple) { + auto test_root_dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(test_root_dir); + std::string test_root = test_root_dir->Str(); + + ASSERT_OK_AND_ASSIGN(auto meta, + WriteGlobalIndex(test_root, data_type_, options_, array_, Range(0, 3))); + ASSERT_OK_AND_ASSIGN(auto reader, + CreateGlobalIndexReader(test_root, data_type_, options_, meta)); + { + // recall all data + ASSERT_OK_AND_ASSIGN( + auto scored_result, + reader->VisitVectorSearch(std::make_shared( + /*field_name=*/"f0", /*limit=*/4, query_, /*filter=*/nullptr, + /*predicate=*/nullptr, /*distance_type=*/std::nullopt, /*options=*/options_))); + CheckResult(scored_result, {3l, 1l, 2l, 0l}, {0.01f, 2.01f, 2.21f, 4.21f}); + } + { + // small limit + ASSERT_OK_AND_ASSIGN( + auto scored_result, + reader->VisitVectorSearch(std::make_shared( + /*field_name=*/"f0", /*limit=*/3, query_, /*filter=*/nullptr, + /*predicate=*/nullptr, /*distance_type=*/std::nullopt, /*options=*/options_))); + CheckResult(scored_result, {3l, 1l, 2l}, {0.01f, 2.01f, 2.21f}); + } + { + // visit equal will return all rows + ASSERT_OK_AND_ASSIGN(auto is_null_result, reader->VisitIsNull()); + ASSERT_FALSE(is_null_result); + } +} + +TEST_F(LuminaGlobalIndexTest, TestWithFilter) { + auto test_root_dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(test_root_dir); + std::string test_root = test_root_dir->Str(); + + ASSERT_OK_AND_ASSIGN(auto meta, + WriteGlobalIndex(test_root, data_type_, options_, array_, Range(0, 3))); + ASSERT_OK_AND_ASSIGN(auto reader, + CreateGlobalIndexReader(test_root, data_type_, options_, meta)); + { + ASSERT_OK_AND_ASSIGN( + auto scored_result, + reader->VisitVectorSearch(std::make_shared( + /*field_name=*/"f0", /*limit=*/2, query_, /*filter=*/nullptr, + /*predicate=*/nullptr, /*distance_type=*/std::nullopt, /*options=*/options_))); + CheckResult(scored_result, {3l, 1l}, {0.01f, 2.01f}); + } + { + auto filter = [](int64_t id) -> bool { return id < 3; }; + ASSERT_OK_AND_ASSIGN( + auto scored_result, + reader->VisitVectorSearch(std::make_shared( + /*field_name=*/"f0", /*limit=*/2, query_, filter, + /*predicate=*/nullptr, /*distance_type=*/std::nullopt, /*options=*/options_))); + CheckResult(scored_result, {1l, 2l}, {2.01f, 2.21f}); + } + { + auto filter = [](int64_t id) -> bool { return id < 3; }; + ASSERT_OK_AND_ASSIGN( + auto scored_result, + reader->VisitVectorSearch(std::make_shared( + /*field_name=*/"f0", /*limit=*/4, query_, filter, + /*predicate=*/nullptr, /*distance_type=*/std::nullopt, /*options=*/options_))); + CheckResult(scored_result, {1l, 2l, 0l}, {2.01f, 2.21f, 4.21f}); + } +} + +TEST_F(LuminaGlobalIndexTest, TestInvalidInputs) { + auto test_root_dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(test_root_dir); + std::string index_root = test_root_dir->Str(); + // invalid inputs in write + { + auto data_type = arrow::int32(); + ASSERT_NOK_WITH_MSG(WriteGlobalIndex(index_root, data_type, options_, array_, Range(0, 3)), + "arrow schema must be struct type when create LuminaIndexWriter"); + } + { + auto data_type = arrow::struct_({arrow::field("f1", arrow::list(arrow::float32()))}); + ASSERT_NOK_WITH_MSG(WriteGlobalIndex(index_root, data_type, options_, array_, Range(0, 3)), + "field f0 not exist in arrow schema when create LuminaIndexWriter"); + } + { + auto data_type = arrow::struct_({arrow::field("f0", arrow::float32())}); + ASSERT_NOK_WITH_MSG(WriteGlobalIndex(index_root, data_type, options_, array_, Range(0, 3)), + "field type must be list[float] when create LuminaIndexWriter"); + } + { + auto data_type = arrow::struct_({arrow::field("f0", arrow::list(arrow::float64()))}); + ASSERT_NOK_WITH_MSG(WriteGlobalIndex(index_root, data_type, options_, array_, Range(0, 3)), + "field type must be list[float] when create LuminaIndexWriter"); + } + { + std::shared_ptr array = arrow::ipc::internal::json::ArrayFromJSON(data_type_, + R"([ + [[0.0, 0.0, 0.0, 0.0]], + null + ])") + .ValueOrDie(); + ASSERT_NOK_WITH_MSG(WriteGlobalIndex(index_root, data_type_, options_, array, Range(0, 2)), + "arrow_array in LuminaIndexWriter is invalid, must not null"); + } + { + std::shared_ptr array = arrow::ipc::internal::json::ArrayFromJSON(data_type_, + R"([ + [[0.0, 0.0, 0.0, 0.0]], + [[0.0, 1.0, 0.0, null]] + ])") + .ValueOrDie(); + ASSERT_NOK_WITH_MSG(WriteGlobalIndex(index_root, data_type_, options_, array, Range(0, 2)), + "field value array in LuminaIndexWriter is invalid, must not null"); + } + { + std::shared_ptr array = arrow::ipc::internal::json::ArrayFromJSON(data_type_, + R"([ + [[0.0, 0.0, 0.0, 0.0]], + [[0.0, 1.0, 0.0]] + ])") + .ValueOrDie(); + ASSERT_NOK_WITH_MSG( + WriteGlobalIndex(index_root, data_type_, options_, array, Range(0, 2)), + "invalid input array in LuminaIndexWriter, length of field array [2] multiplied " + "dimension [4] must match length of field value array [7]"); + } + + { + // invalid inputs in read + auto test_root_dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(test_root_dir); + std::string index_root = test_root_dir->Str(); + ASSERT_OK_AND_ASSIGN( + auto meta, WriteGlobalIndex(index_root, data_type_, options_, array_, Range(0, 3))); + // read + { + auto fake_meta = meta; + fake_meta.metadata = nullptr; + ASSERT_NOK_WITH_MSG( + CreateGlobalIndexReader(index_root, data_type_, options_, /*meta=*/fake_meta), + "Lumina global index must have meta data"); + } + { + auto fake_meta = meta; + auto fake_index_meta_json = StringUtils::Replace( + std::string(fake_meta.metadata->data(), fake_meta.metadata->size()), + /*search_string=*/"l2", /*replacement=*/"unknown"); + fake_meta.metadata = std::make_shared(fake_index_meta_json, pool_.get()); + ASSERT_NOK_WITH_MSG( + CreateGlobalIndexReader(index_root, data_type_, options_, fake_meta), + "invalid distance type unknown for lumina"); + } + { + auto global_index = std::make_shared(options_); + auto path_factory = std::make_shared(index_root); + auto file_reader = std::make_shared(fs_, path_factory); + + ASSERT_NOK_WITH_MSG(global_index->CreateReader(CreateArrowSchema(data_type_).get(), + file_reader, {meta, meta}, pool_), + "lumina index only has one index file per shard"); + } + { + auto data_type = arrow::struct_({arrow::field("f0", arrow::list(arrow::float32())), + arrow::field("f1", arrow::list(arrow::float32()))}); + ASSERT_NOK_WITH_MSG(CreateGlobalIndexReader(index_root, data_type, options_, meta), + "LuminaGlobalIndex now only support one field"); + } + { + auto data_type = arrow::struct_({arrow::field("f0", arrow::float32())}); + ASSERT_NOK_WITH_MSG(CreateGlobalIndexReader(index_root, data_type, options_, meta), + "field type must be list[float] when create LuminaIndexReader"); + } + { + auto data_type = arrow::struct_({arrow::field("f0", arrow::list(arrow::float64()))}); + ASSERT_NOK_WITH_MSG(CreateGlobalIndexReader(index_root, data_type, options_, meta), + "field type must be list[float] when create LuminaIndexReader"); + } + { + auto fake_meta = meta; + fake_meta.file_path = "non-exist-file"; + ASSERT_NOK_WITH_MSG( + CreateGlobalIndexReader(index_root, data_type_, options_, fake_meta), + "non-exist-file\' not exists"); + } + { + ASSERT_OK_AND_ASSIGN(auto reader, + CreateGlobalIndexReader(index_root, data_type_, options_, meta)); + ASSERT_NOK_WITH_MSG(reader->VisitVectorSearch(std::make_shared( + "f1", + /*limit=*/2, query_, /*filter=*/nullptr, + PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"f0", + FieldType::BIGINT, Literal(5l)), + /*distance_type=*/std::nullopt, + /*options=*/std::map())), + "lumina index not support predicate in VisitVectorSearch"); + } + { + ASSERT_OK_AND_ASSIGN(auto reader, + CreateGlobalIndexReader(index_root, data_type_, options_, meta)); + ASSERT_NOK_WITH_MSG(reader->VisitVectorSearch(std::make_shared( + "f1", + /*limit=*/2, query_, /*filter=*/nullptr, + /*predicate=*/nullptr, + /*distance_type=*/VectorSearch::DistanceType::COSINE, + /*options=*/std::map())), + "distance type for index and search not match"); + } + { + ASSERT_OK_AND_ASSIGN(auto reader, + CreateGlobalIndexReader(index_root, data_type_, options_, meta)); + auto query = query_; + query.push_back(1.0f); + ASSERT_NOK_WITH_MSG(reader->VisitVectorSearch(std::make_shared( + "f1", + /*limit=*/2, query, /*filter=*/nullptr, + /*predicate=*/nullptr, + /*distance_type=*/std::nullopt, + /*options=*/std::map())), + "dimension for index and search not match"); + } + { + ASSERT_OK_AND_ASSIGN(auto reader, + CreateGlobalIndexReader(index_root, data_type_, options_, meta)); + auto fake_options = options_; + fake_options["lumina.index.type"] = "diskann"; + ASSERT_NOK_WITH_MSG(reader->VisitVectorSearch(std::make_shared( + "f1", + /*limit=*/2, query_, /*filter=*/nullptr, + /*predicate=*/nullptr, + /*distance_type=*/std::nullopt, + /*options=*/fake_options)), + "index type for index and search not match"); + } + } +} + +TEST_F(LuminaGlobalIndexTest, TestHighCardinalityAndMultiThreadSearch) { + int64_t seed = DateTimeUtils::GetCurrentUTCTimeUs(); + std::srand(seed); + auto test_root_dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(test_root_dir); + std::string test_root = test_root_dir->Str(); + + auto array = CreateRandomVector(/*element_size*/ 10000, /*dimension=*/4); + ASSERT_OK_AND_ASSIGN( + auto meta, WriteGlobalIndex(test_root, data_type_, options_, array, Range(0, 10000 - 1))); + ASSERT_OK_AND_ASSIGN(auto reader, + CreateGlobalIndexReader(test_root, data_type_, options_, meta)); + + auto search_with_filter = [&]() { + int32_t limit = paimon::test::RandomNumber(1, 100); + auto filter = [](int64_t id) -> bool { return id % 2; }; + ASSERT_OK_AND_ASSIGN( + auto scored_result, + reader->VisitVectorSearch(std::make_shared( + "f0", limit, query_, filter, + /*predicate=*/nullptr, /*distance_type=*/std::nullopt, /*options=*/options_))); + auto typed_result = std::dynamic_pointer_cast(scored_result); + ASSERT_TRUE(typed_result); + ASSERT_EQ(typed_result->bitmap_.Cardinality(), limit); + }; + + auto search = [&]() { + int32_t limit = paimon::test::RandomNumber(1, 100); + ASSERT_OK_AND_ASSIGN( + auto scored_result, + reader->VisitVectorSearch(std::make_shared( + "f0", limit, query_, /*filter=*/nullptr, + /*predicate=*/nullptr, /*distance_type=*/std::nullopt, /*options=*/options_))); + auto typed_result = std::dynamic_pointer_cast(scored_result); + ASSERT_TRUE(typed_result); + ASSERT_EQ(typed_result->bitmap_.Cardinality(), limit); + }; + + std::vector threads; + for (int32_t i = 0; i < 5; ++i) { + threads.emplace_back(search); + } + for (int32_t i = 0; i < 5; ++i) { + threads.emplace_back(search_with_filter); + } + + for (auto& t : threads) { + if (t.joinable()) { + t.join(); + } + } +} + +} // namespace paimon::lumina::test diff --git a/src/paimon/global_index/lumina/lumina_memory_pool.h b/src/paimon/global_index/lumina/lumina_memory_pool.h new file mode 100644 index 00000000..e5dcee9d --- /dev/null +++ b/src/paimon/global_index/lumina/lumina_memory_pool.h @@ -0,0 +1,59 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +#include "paimon/memory/memory_pool.h" +namespace paimon::lumina { +class LuminaMemoryPool : public std::pmr::memory_resource { + public: + explicit LuminaMemoryPool(const std::shared_ptr& pool) : pool_(pool) {} + + ~LuminaMemoryPool() override = default; + + LuminaMemoryPool(const LuminaMemoryPool&) = delete; + LuminaMemoryPool& operator=(const LuminaMemoryPool&) = delete; + + std::shared_ptr GetPaimonPool() const { + return pool_; + } + + private: + void* do_allocate(std::size_t bytes, std::size_t alignment) override { + static const size_t min_alignment = alignof(std::max_align_t); + if (alignment < min_alignment) { + alignment = min_alignment; + } + return pool_->Malloc(bytes, alignment); + } + + void do_deallocate(void* p, std::size_t bytes, std::size_t alignment) override { + pool_->Free(p, bytes, alignment); + } + + bool do_is_equal(const std::pmr::memory_resource& other) const noexcept override { + return this == &other; + } + + private: + std::shared_ptr pool_; +}; +} // namespace paimon::lumina diff --git a/src/paimon/global_index/lumina/lumina_utils.h b/src/paimon/global_index/lumina/lumina_utils.h new file mode 100644 index 00000000..60f79785 --- /dev/null +++ b/src/paimon/global_index/lumina/lumina_utils.h @@ -0,0 +1,94 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include "lumina/core/Status.h" +#include "paimon/status.h" + +namespace paimon::lumina { +#define PAIMON_RETURN_NOT_OK_FROM_LUMINA(LUMINA_STATUS) \ + do { \ + ::lumina::core::Status __s = (LUMINA_STATUS); \ + if (PAIMON_UNLIKELY(!(__s).IsOk())) { \ + return LuminaToPaimonStatus(__s); \ + } \ + } while (false) + +#define PAIMON_ASSIGN_OR_RAISE_IMPL_FROM_LUMINA(result_name, lhs, rexpr) \ + auto&& result_name = (rexpr); \ + PAIMON_RETURN_IF_(!(result_name).IsOk(), LuminaToPaimonStatus((result_name).GetStatus()), \ + PAIMON_STRINGIFY(rexpr)); \ + lhs = std::move(result_name).TakeValue(); + +#define PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA(lhs, rexpr) \ + PAIMON_ASSIGN_OR_RAISE_IMPL_FROM_LUMINA( \ + PAIMON_ASSIGN_OR_RAISE_NAME(_error_or_value, __COUNTER__), lhs, (rexpr)); + +inline ::lumina::core::Status PaimonToLuminaStatus(const Status& status) { + switch (status.code()) { + case StatusCode::OK: + return ::lumina::core::Status::Ok(); + case StatusCode::OutOfMemory: + return ::lumina::core::Status(::lumina::core::ErrorCode::OutOfMemory, status.message()); + case StatusCode::IOError: + return ::lumina::core::Status(::lumina::core::ErrorCode::IoError, status.message()); + case StatusCode::NotImplemented: + return ::lumina::core::Status(::lumina::core::ErrorCode::NotSupported, + status.message()); + case StatusCode::NotExist: + return ::lumina::core::Status(::lumina::core::ErrorCode::NotFound, status.message()); + case StatusCode::Exist: + return ::lumina::core::Status(::lumina::core::ErrorCode::AlreadyExists, + status.message()); + default: + return ::lumina::core::Status(::lumina::core::ErrorCode::InvalidArgument, + status.message()); + } +} + +inline Status LuminaToPaimonStatus(const ::lumina::core::Status& status) { + switch (status.Code()) { + case ::lumina::core::ErrorCode::Ok: + return Status::OK(); + case ::lumina::core::ErrorCode::NotFound: + return Status::NotExist(status.Message()); + case ::lumina::core::ErrorCode::AlreadyExists: + return Status::Exist(status.Message()); + case ::lumina::core::ErrorCode::NotSupported: + return Status::NotImplemented(status.Message()); + case ::lumina::core::ErrorCode::IoError: + return Status::IOError(status.Message()); + case ::lumina::core::ErrorCode::OutOfMemory: + return Status::OutOfMemory(status.Message()); + default: + return Status::Invalid(status.Message()); + } +} + +class LuminaDefines { + public: + LuminaDefines() = delete; + ~LuminaDefines() = delete; + static constexpr char kOptionKeyPrefix[] = "lumina."; + static constexpr char kIdentifier[] = "lumina"; +}; + +} // namespace paimon::lumina diff --git a/third_party/versions.txt b/third_party/versions.txt index c2759999..be93f35b 100644 --- a/third_party/versions.txt +++ b/third_party/versions.txt @@ -84,6 +84,10 @@ PAIMON_RAPIDJSON_BUILD_VERSION=232389d4f1012dddec4ef84861face2d2ba85709 PAIMON_RAPIDJSON_BUILD_SHA256_CHECKSUM=b9290a9a6d444c8e049bd589ab804e0ccf2b05dc5984a19ed5ae75d090064806 PAIMON_RAPIDJSON_PKG_NAME=rapidjson-${PAIMON_RAPIDJSON_BUILD_VERSION}.tar.gz +PAIMON_LUMINA_BUILD_VERSION=0.2.3 +PAIMON_LUMINA_BUILD_SHA256_CHECKSUM=49de9548a8b81f6a5e3794e8e65fb32023a1c4bdb019028ce55cbcc98c12be97 +PAIMON_LUMINA_PKG_NAME=lumina_release-${PAIMON_LUMINA_BUILD_VERSION}.tar.gz + PAIMON_JINDOSDK_C_BUILD_VERSION=6.10.2 PAIMON_JINDOSDK_C_LINUX_X86_64_BUILD_SHA256_CHECKSUM=23e61c9815fab1cd88c369445bdbe1eab02cc09bafed3bb5118ecaf5b2fbc518 PAIMON_JINDOSDK_C_LINUX_X86_64_PKG_NAME=jindosdk-${PAIMON_JINDOSDK_C_BUILD_VERSION}-linux.tar.gz @@ -106,7 +110,6 @@ PAIMON_JIEBA_BUILD_VERSION=v5.6.0 PAIMON_JIEBA_BUILD_SHA256_CHECKSUM=e6e517b778e0f4a99cbed1ee3eaa041616b74bc685e03a6ca08887ad9cedfe49 PAIMON_JIEBA_PKG_NAME=jieba-${PAIMON_JIEBA_BUILD_VERSION}.tar.gz -# Boost source package is bundled at third_party/boost/ PAIMON_BOOST_BUILD_VERSION=1_66_0 PAIMON_BOOST_BUILD_SHA256_CHECKSUM=28e9200637800fbfd1292b2c6876189dba7e8e1c5282c71fac6515e96f7af2b0 PAIMON_BOOST_PKG_NAME=boost_${PAIMON_BOOST_BUILD_VERSION}.tar.gz @@ -130,6 +133,7 @@ DEPENDENCIES=( "PAIMON_GLOG_URL ${PAIMON_GLOG_PKG_NAME} ${THIRDPARTY_MIRROR_URL}https://github.com/google/glog/archive/${PAIMON_GLOG_BUILD_VERSION}.tar.gz" "PAIMON_RAPIDJSON_URL ${PAIMON_RAPIDJSON_PKG_NAME} ${THIRDPARTY_MIRROR_URL}https://github.com/miloyip/rapidjson/archive/${PAIMON_RAPIDJSON_BUILD_VERSION}.tar.gz" "PAIMON_RE2_URL ${PAIMON_RE2_PKG_NAME} ${THIRDPARTY_MIRROR_URL}https://github.com/google/re2/archive/${PAIMON_RE2_BUILD_VERSION}.tar.gz" + "PAIMON_LUMINA_URL ${PAIMON_LUMINA_PKG_NAME} https://paimon-cpp.oss-cn-beijing.aliyuncs.com/thirdparty/lumina/lumina_release-${PAIMON_LUMINA_BUILD_VERSION}.tar.gz" "PAIMON_JINDOSDK_C_LINUX_X86_64_URL ${PAIMON_JINDOSDK_C_LINUX_X86_64_PKG_NAME} https://jindodata-binary.oss-cn-shanghai.aliyuncs.com/release/${PAIMON_JINDOSDK_C_BUILD_VERSION}/jindosdk-${PAIMON_JINDOSDK_C_BUILD_VERSION}-linux.tar.gz" "PAIMON_JINDOSDK_C_LINUX_AARCH64_URL ${PAIMON_JINDOSDK_C_LINUX_AARCH64_PKG_NAME} https://jindodata-binary.oss-cn-shanghai.aliyuncs.com/release/${PAIMON_JINDOSDK_C_BUILD_VERSION}/jindosdk-${PAIMON_JINDOSDK_C_BUILD_VERSION}-linux-el7-aarch64.tar.gz" "PAIMON_JINDOSDK_C_MACOS_X86_64_URL ${PAIMON_JINDOSDK_C_MACOS_X86_64_PKG_NAME} https://jindodata-binary.oss-cn-shanghai.aliyuncs.com/release/${PAIMON_JINDOSDK_C_BUILD_VERSION}/jindosdk-${PAIMON_JINDOSDK_C_BUILD_VERSION}-macos-11_0-x86_64.tar.gz" @@ -137,4 +141,5 @@ DEPENDENCIES=( "PAIMON_LUCENE_URL ${PAIMON_LUCENE_PKG_NAME} ${THIRDPARTY_MIRROR_URL}https://github.com/luceneplusplus/LucenePlusPlus/archive/refs/tags/rel_${PAIMON_LUCENE_BUILD_VERSION}.tar.gz" "PAIMON_LIMONP_URL ${PAIMON_LIMONP_PKG_NAME} ${THIRDPARTY_MIRROR_URL}https://github.com/yanyiwu/limonp/archive/refs/tags/v${PAIMON_LIMONP_BUILD_VERSION}.tar.gz" "PAIMON_JIEBA_URL ${PAIMON_JIEBA_PKG_NAME} ${THIRDPARTY_MIRROR_URL}https://github.com/yanyiwu/cppjieba/archive/refs/tags/${PAIMON_JIEBA_BUILD_VERSION}.tar.gz" + "PAIMON_BOOST_URL ${PAIMON_BOOST_PKG_NAME} https://paimon-cpp.oss-cn-beijing.aliyuncs.com/thirdparty/boost/${PAIMON_BOOST_PKG_NAME}" )