diff --git a/.github/cmake/subproject-integration/CMakeLists.txt b/.github/cmake/subproject-integration/CMakeLists.txt index dbae74c..a6cb349 100644 --- a/.github/cmake/subproject-integration/CMakeLists.txt +++ b/.github/cmake/subproject-integration/CMakeLists.txt @@ -9,30 +9,18 @@ if(NOT ZVEC_SOURCE_DIR) message(FATAL_ERROR "ZVEC_SOURCE_DIR is required") endif() -option(ENABLE_ZVEC_CORE "Validate zvec-core-only subproject integration" OFF) - -if(ENABLE_ZVEC_CORE) - set(BUILD_ZVEC_CORE_ONLY ON CACHE BOOL "" FORCE) -endif() - add_subdirectory("${ZVEC_SOURCE_DIR}" zvec) set(CMAKE_RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/bin") -if(ENABLE_ZVEC_CORE) - add_executable(zvec_embed_core_only_example - "${ZVEC_SOURCE_DIR}/examples/c++/core/main.cc") - target_link_libraries(zvec_embed_core_only_example PRIVATE zvec_core_shared) -else() - add_executable(zvec_embed_db_example - "${ZVEC_SOURCE_DIR}/examples/c++/db/main.cc") - target_link_libraries(zvec_embed_db_example PRIVATE zvec_shared) +add_executable(zvec_embed_db_example + "${ZVEC_SOURCE_DIR}/examples/c++/db/main.cc") +target_link_libraries(zvec_embed_db_example PRIVATE zvec_shared) - add_executable(zvec_embed_core_example - "${ZVEC_SOURCE_DIR}/examples/c++/core/main.cc") - target_link_libraries(zvec_embed_core_example PRIVATE zvec_shared) +add_executable(zvec_embed_core_example + "${ZVEC_SOURCE_DIR}/examples/c++/core/main.cc") +target_link_libraries(zvec_embed_core_example PRIVATE zvec_core_shared) - add_executable(zvec_embed_ailego_example - "${ZVEC_SOURCE_DIR}/examples/c++/ailego/main.cc") - target_link_libraries(zvec_embed_ailego_example PRIVATE zvec_shared) -endif() +add_executable(zvec_embed_ailego_example + "${ZVEC_SOURCE_DIR}/examples/c++/ailego/main.cc") +target_link_libraries(zvec_embed_ailego_example PRIVATE zvec_ailego_shared) diff --git a/.github/workflows/05-windows-build.yml b/.github/workflows/05-windows-build.yml index 7ed5503..c7f7fc9 100644 --- a/.github/workflows/05-windows-build.yml +++ b/.github/workflows/05-windows-build.yml @@ -120,47 +120,29 @@ jobs: -DCMAKE_CXX_COMPILER_LAUNCHER=sccache cmake --build . --config Release --parallel $env:NPROC - # Copy zvec.dll next to the example executables so Windows can find it. + # Copy zvec DLLs next to the example executables so Windows can find them. # CMake places DLLs in bin/ (RUNTIME output) and import libs in lib/. $buildDir = "$env:GITHUB_WORKSPACE\build" - $found = $false - foreach ($sub in @("$buildDir\bin", "$buildDir\bin\Release", "$buildDir\lib", "$buildDir\lib\Release")) { - if (Test-Path "$sub\zvec.dll") { - Copy-Item "$sub\zvec.dll" -Destination . -Force - Write-Host "Copied zvec.dll from $sub" - $found = $true - break + foreach ($dllName in @("zvec.dll", "zvec_core.dll", "zvec_ailego.dll")) { + $found = $false + foreach ($sub in @("$buildDir\bin", "$buildDir\bin\Release", "$buildDir\lib", "$buildDir\lib\Release")) { + $dllPath = Join-Path $sub $dllName + if (Test-Path $dllPath) { + Copy-Item $dllPath -Destination . -Force + Write-Host "Copied $dllName from $sub" + $found = $true + break + } } - } - if (-not $found) { - Write-Host "WARNING: zvec.dll not found, searching recursively..." - $dll = Get-ChildItem -Path $buildDir -Filter "zvec.dll" -Recurse -ErrorAction SilentlyContinue | Select-Object -First 1 - if ($dll) { - Copy-Item $dll.FullName -Destination . -Force - Write-Host "Copied zvec.dll from $($dll.DirectoryName)" - } else { - Write-Error "zvec.dll not found anywhere under $buildDir" - } - } - - # core-example links libzvec_core directly, so copy zvec_core.dll too. - $found = $false - foreach ($sub in @("$buildDir\bin", "$buildDir\bin\Release", "$buildDir\lib", "$buildDir\lib\Release")) { - if (Test-Path "$sub\zvec_core.dll") { - Copy-Item "$sub\zvec_core.dll" -Destination . -Force - Write-Host "Copied zvec_core.dll from $sub" - $found = $true - break - } - } - if (-not $found) { - Write-Host "WARNING: zvec_core.dll not found, searching recursively..." - $dll = Get-ChildItem -Path $buildDir -Filter "zvec_core.dll" -Recurse -ErrorAction SilentlyContinue | Select-Object -First 1 - if ($dll) { - Copy-Item $dll.FullName -Destination . -Force - Write-Host "Copied zvec_core.dll from $($dll.DirectoryName)" - } else { - Write-Error "zvec_core.dll not found anywhere under $buildDir" + if (-not $found) { + Write-Host "WARNING: $dllName not found, searching recursively..." + $dll = Get-ChildItem -Path $buildDir -Filter $dllName -Recurse -ErrorAction SilentlyContinue | Select-Object -First 1 + if ($dll) { + Copy-Item $dll.FullName -Destination . -Force + Write-Host "Copied $dllName from $($dll.DirectoryName)" + } else { + Write-Error "$dllName not found anywhere under $buildDir" + } } } diff --git a/.github/workflows/08-cmake-subproject-integration.yml b/.github/workflows/08-cmake-subproject-integration.yml index 5709f4d..7cb2dd1 100644 --- a/.github/workflows/08-cmake-subproject-integration.yml +++ b/.github/workflows/08-cmake-subproject-integration.yml @@ -11,19 +11,9 @@ permissions: jobs: subproject-integration: - name: Subproject Build & Example (${{ matrix.mode }}) + name: Subproject Build & Example runs-on: ubuntu-24.04 timeout-minutes: 90 - strategy: - fail-fast: false - matrix: - include: - - mode: linux-x64 - cache_key: full - enable_zvec_core: false - - mode: linux-x64-enable-zvec-core - cache_key: core-only - enable_zvec_core: true steps: - name: Checkout code @@ -34,7 +24,7 @@ jobs: - name: Setup ccache uses: hendrikmuhs/ccache-action@v1.2 with: - key: cmake-subproject-linux-x64-${{ matrix.cache_key }} + key: cmake-subproject-linux-x64 max-size: 150M - name: Set up Python @@ -68,8 +58,8 @@ jobs: - name: Configure subproject integration run: | - SMOKE_SOURCE="$RUNNER_TEMP/zvec-subproject-nightly-${{ matrix.cache_key }}" - SMOKE_BUILD="$RUNNER_TEMP/zvec-subproject-nightly-build-${{ matrix.cache_key }}" + SMOKE_SOURCE="$RUNNER_TEMP/zvec-subproject-nightly" + SMOKE_BUILD="$RUNNER_TEMP/zvec-subproject-nightly-build" mkdir -p "$SMOKE_SOURCE" "$SMOKE_BUILD" cp "$GITHUB_WORKSPACE/.github/cmake/subproject-integration/CMakeLists.txt" \ @@ -79,20 +69,14 @@ jobs: -DZVEC_SOURCE_DIR="$GITHUB_WORKSPACE" -DCMAKE_BUILD_TYPE=Release -DENABLE_WERROR=ON + -DUSE_OSS_MIRROR=ON -DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + -DBUILD_TOOLS=OFF + -DBUILD_C_BINDINGS=OFF + -DBUILD_PYTHON_BINDINGS=OFF ) - if [[ "${{ matrix.enable_zvec_core }}" == "true" ]]; then - CMAKE_ARGS+=(-DENABLE_ZVEC_CORE=ON) - else - CMAKE_ARGS+=( - -DBUILD_TOOLS=OFF - -DBUILD_C_BINDINGS=OFF - -DBUILD_PYTHON_BINDINGS=OFF - ) - fi - cmake -S "$SMOKE_SOURCE" -B "$SMOKE_BUILD" -G Ninja \ "${CMAKE_ARGS[@]}" @@ -100,7 +84,6 @@ jobs: shell: bash - name: Build subproject targets - if: ${{ ! matrix.enable_zvec_core }} run: | cmake --build "$SMOKE_BUILD" --target \ core_knn_diskann \ @@ -111,32 +94,9 @@ jobs: shell: bash - name: Run subproject examples - if: ${{ ! matrix.enable_zvec_core }} run: | cd "$SMOKE_BUILD" ./bin/zvec_embed_db_example ./bin/zvec_embed_core_example ./bin/zvec_embed_ailego_example shell: bash - - - name: Build enable_zvec_core targets - if: ${{ matrix.enable_zvec_core }} - run: | - cmake --build "$SMOKE_BUILD" --target \ - zvec_core_shared \ - zvec_embed_core_only_example \ - --parallel "$NPROC" - shell: bash - - - name: Run enable_zvec_core example - if: ${{ matrix.enable_zvec_core }} - run: | - cd "$SMOKE_BUILD" - ./bin/zvec_embed_core_only_example - ldd ./bin/zvec_embed_core_only_example | tee zvec-core-only-ldd.txt - grep -q 'libzvec_core\.so' zvec-core-only-ldd.txt - if grep -q 'libzvec\.so' zvec-core-only-ldd.txt; then - echo "core-only example unexpectedly depends on libzvec.so" - exit 1 - fi - shell: bash diff --git a/CMakeLists.txt b/CMakeLists.txt index 392dbda..f9ad6bc 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -74,25 +74,16 @@ endif() include_directories(${PROJECT_ROOT_DIR}/src/include) include_directories(${PROJECT_ROOT_DIR}/src) -option(BUILD_ZVEC_CORE_ONLY "Build only zvec-core and its all-in-one shared library" OFF) -message(STATUS "BUILD_ZVEC_CORE_ONLY:${BUILD_ZVEC_CORE_ONLY}") - option(BUILD_ZVEC_SHARED "Build all-in-one C++ shared library libzvec" ON) +option(BUILD_ZVEC_AILEGO_SHARED "Build all-in-one zvec-ailego shared library libzvec_ailego" ON) option(BUILD_ZVEC_CORE_SHARED "Build all-in-one zvec-core shared library libzvec_core" ON) option(BUILD_PYTHON_BINDINGS "Build Python bindings using pybind11" OFF) option(BUILD_C_BINDINGS "Build C bindings" ON) option(BUILD_TOOLS "Build tools" ON) -if(BUILD_ZVEC_CORE_ONLY) - set(BUILD_ZVEC_SHARED OFF CACHE BOOL "Build all-in-one C++ shared library libzvec" FORCE) - set(BUILD_ZVEC_CORE_SHARED ON CACHE BOOL "Build all-in-one zvec-core shared library libzvec_core" FORCE) - set(BUILD_PYTHON_BINDINGS OFF CACHE BOOL "Build Python bindings using pybind11" FORCE) - set(BUILD_C_BINDINGS OFF CACHE BOOL "Build C bindings" FORCE) - set(BUILD_TOOLS OFF CACHE BOOL "Build tools" FORCE) -endif() - message(STATUS "BUILD_ZVEC_SHARED:${BUILD_ZVEC_SHARED}") +message(STATUS "BUILD_ZVEC_AILEGO_SHARED:${BUILD_ZVEC_AILEGO_SHARED}") message(STATUS "BUILD_ZVEC_CORE_SHARED:${BUILD_ZVEC_CORE_SHARED}") message(STATUS "BUILD_PYTHON_BINDINGS:${BUILD_PYTHON_BINDINGS}") message(STATUS "BUILD_C_BINDINGS:${BUILD_C_BINDINGS}") @@ -151,17 +142,9 @@ message(STATUS "USE_OSS_MIRROR:${USE_OSS_MIRROR}") cc_directory(thirdparty) cc_directories(src) -if(BUILD_ZVEC_CORE_ONLY) - # Keep the include context from tests/CMakeLists.txt for core tests that include tests/test_util.h. - include_directories(${PROJECT_ROOT_DIR}) - cc_directory(tests/core) -else() - cc_directories(tests) -endif() +cc_directories(tests) -if(NOT BUILD_ZVEC_CORE_ONLY) - add_custom_target(clang_tidy_deps DEPENDS zvec_proto ARROW.BUILD glog gflags Lz4.BUILD) -endif() +add_custom_target(clang_tidy_deps DEPENDS zvec_proto ARROW.BUILD glog gflags Lz4.BUILD) if(BUILD_TOOLS) cc_directories(tools) diff --git a/examples/c++/CMakeLists.txt b/examples/c++/CMakeLists.txt index dc59961..1de9029 100644 --- a/examples/c++/CMakeLists.txt +++ b/examples/c++/CMakeLists.txt @@ -16,10 +16,6 @@ endif() get_filename_component(ZVEC_ROOT_DIR "${CMAKE_CURRENT_LIST_DIR}/../.." ABSOLUTE) set(ZVEC_INCLUDE_DIR ${ZVEC_ROOT_DIR}/src/include) set(ZVEC_LIB_DIR ${ZVEC_ROOT_DIR}/${HOST_BUILD_DIR}/lib) -set(ZVEC_LINK_VIA_CMAKE_TARGET OFF) -if(TARGET zvec_shared OR TARGET zvec_core_shared) - set(ZVEC_LINK_VIA_CMAKE_TARGET ON) -endif() # Add include and library search paths include_directories(${ZVEC_INCLUDE_DIR}) @@ -32,10 +28,6 @@ if(CMAKE_BUILD_TYPE) list(APPEND ZVEC_LIB_SEARCH_DIRS ${ZVEC_CONFIG_LIB_DIR}) endif() endif() -if(NOT ZVEC_LINK_VIA_CMAKE_TARGET) - link_directories(${ZVEC_LIB_SEARCH_DIRS}) - set(CMAKE_BUILD_RPATH ${ZVEC_LIB_SEARCH_DIRS}) -endif() if(WIN32) set(CMAKE_MSVC_RUNTIME_LIBRARY "MultiThreaded$<$:Debug>") endif() @@ -47,6 +39,7 @@ function(zvec_find_shared_library OUT_VAR LIB_NAME) NAMES ${LIB_NAME}_shared ${LIB_NAME} PATHS ${ZVEC_LIB_SEARCH_DIRS} NO_DEFAULT_PATH + NO_CMAKE_FIND_ROOT_PATH ) else() set(ZVEC_ORIGINAL_LIBRARY_SUFFIXES ${CMAKE_FIND_LIBRARY_SUFFIXES}) @@ -59,90 +52,61 @@ function(zvec_find_shared_library OUT_VAR LIB_NAME) NAMES ${LIB_NAME} PATHS ${ZVEC_LIB_SEARCH_DIRS} NO_DEFAULT_PATH + NO_CMAKE_FIND_ROOT_PATH ) set(CMAKE_FIND_LIBRARY_SUFFIXES "${ZVEC_ORIGINAL_LIBRARY_SUFFIXES}") endif() set(${OUT_VAR} "${${OUT_VAR}}" PARENT_SCOPE) endfunction() -set(ZVEC_CORE_EXAMPLE_DEFAULT ON) -if(NOT TARGET zvec_core_shared) - zvec_find_shared_library(ZVEC_CORE_SHARED_LIBRARY zvec_core) - if(NOT ZVEC_CORE_SHARED_LIBRARY) - set(ZVEC_CORE_EXAMPLE_DEFAULT OFF) +function(zvec_require_shared_library OUT_VAR LIB_NAME) + zvec_find_shared_library(${OUT_VAR} ${LIB_NAME}) + if(NOT ${OUT_VAR}) + message(FATAL_ERROR + "lib${LIB_NAME} shared library was not found in ${ZVEC_LIB_SEARCH_DIRS}. " + "Build zvec first, or pass -DHOST_BUILD_DIR=.") endif() -endif() + set(${OUT_VAR} "${${OUT_VAR}}" PARENT_SCOPE) +endfunction() -set(ZVEC_FULL_EXAMPLES_DEFAULT ON) -if(NOT TARGET zvec_shared) - zvec_find_shared_library(ZVEC_SHARED_LIBRARY zvec) - if(NOT ZVEC_SHARED_LIBRARY) - set(ZVEC_FULL_EXAMPLES_DEFAULT OFF) - endif() -endif() - -option(BUILD_ZVEC_CORE_EXAMPLE "Build zvec-core example linked only with libzvec_core" ${ZVEC_CORE_EXAMPLE_DEFAULT}) -option(BUILD_ZVEC_FULL_EXAMPLES "Build examples linked with libzvec" ${ZVEC_FULL_EXAMPLES_DEFAULT}) +zvec_require_shared_library(ZVEC_SHARED_LIBRARY zvec) +zvec_require_shared_library(ZVEC_AILEGO_SHARED_LIBRARY zvec_ailego) +zvec_require_shared_library(ZVEC_CORE_SHARED_LIBRARY zvec_core) # --- Create INTERFACE target for libzvec (all-in-one C++ shared library) --- # libzvec.so/.dylib/.dll already bundles all zvec internal components # (zvec, zvec_core, zvec_ailego, zvec_turbo), so no individual dependency # libraries need to be specified by the consumer. -if(BUILD_ZVEC_FULL_EXAMPLES) - add_library(zvec-lib INTERFACE) - if(TARGET zvec_shared) - target_link_libraries(zvec-lib INTERFACE zvec_shared) - else() - if(ZVEC_SHARED_LIBRARY) - target_link_libraries(zvec-lib INTERFACE "${ZVEC_SHARED_LIBRARY}") - else() - message(FATAL_ERROR - "libzvec shared library was not found in ${ZVEC_LIB_SEARCH_DIRS}. " - "Build zvec with -DBUILD_ZVEC_SHARED=ON or disable BUILD_ZVEC_FULL_EXAMPLES.") - endif() - endif() -endif() +add_library(zvec-lib INTERFACE) +target_link_libraries(zvec-lib INTERFACE "${ZVEC_SHARED_LIBRARY}") + +# --- Create INTERFACE target for libzvec_ailego (ailego-only all-in-one library) --- +# The ailego example intentionally depends only on libzvec_ailego. +add_library(zvec-ailego-lib INTERFACE) +target_link_libraries(zvec-ailego-lib INTERFACE "${ZVEC_AILEGO_SHARED_LIBRARY}") # --- Create INTERFACE target for libzvec_core (core-only all-in-one library) --- # The core example intentionally depends only on libzvec_core. -if(BUILD_ZVEC_CORE_EXAMPLE) - add_library(zvec-core-lib INTERFACE) - if(TARGET zvec_core_shared) - target_link_libraries(zvec-core-lib INTERFACE zvec_core_shared) - else() - zvec_find_shared_library(ZVEC_CORE_SHARED_LIBRARY zvec_core) - if(ZVEC_CORE_SHARED_LIBRARY) - target_link_libraries(zvec-core-lib INTERFACE "${ZVEC_CORE_SHARED_LIBRARY}") - else() - message(FATAL_ERROR - "libzvec_core shared library was not found in ${ZVEC_LIB_SEARCH_DIRS}. " - "Build zvec with -DBUILD_ZVEC_CORE_SHARED=ON or " - "-DBUILD_ZVEC_CORE_ONLY=ON, or disable BUILD_ZVEC_CORE_EXAMPLE.") - endif() - endif() -endif() +add_library(zvec-core-lib INTERFACE) +target_link_libraries(zvec-core-lib INTERFACE "${ZVEC_CORE_SHARED_LIBRARY}") # --- Executables --- set(ZVEC_EXAMPLE_TARGETS) -if(BUILD_ZVEC_FULL_EXAMPLES) - add_executable(db-example db/main.cc) - target_link_libraries(db-example PRIVATE zvec-lib) - if(ANDROID) - target_link_libraries(db-example PRIVATE log) - endif() - list(APPEND ZVEC_EXAMPLE_TARGETS db-example) - - add_executable(ailego-example ailego/main.cc) - target_link_libraries(ailego-example PRIVATE zvec-lib) - list(APPEND ZVEC_EXAMPLE_TARGETS ailego-example) +add_executable(db-example db/main.cc) +target_link_libraries(db-example PRIVATE zvec-lib) +if(ANDROID) + target_link_libraries(db-example PRIVATE log) endif() +list(APPEND ZVEC_EXAMPLE_TARGETS db-example) -if(BUILD_ZVEC_CORE_EXAMPLE) - add_executable(core-example core/main.cc) - target_link_libraries(core-example PRIVATE zvec-core-lib) - list(APPEND ZVEC_EXAMPLE_TARGETS core-example) -endif() +add_executable(ailego-example ailego/main.cc) +target_link_libraries(ailego-example PRIVATE zvec-ailego-lib) +list(APPEND ZVEC_EXAMPLE_TARGETS ailego-example) + +add_executable(core-example core/main.cc) +target_link_libraries(core-example PRIVATE zvec-core-lib) +list(APPEND ZVEC_EXAMPLE_TARGETS core-example) # Strip symbols to reduce executable size if(CMAKE_BUILD_TYPE STREQUAL "Release" AND ANDROID) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 9cf11f3..5fc4e34 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -8,11 +8,8 @@ git_version(ZVEC_VERSION ${CMAKE_CURRENT_SOURCE_DIR}) cc_directory(ailego) cc_directory(turbo) cc_directory(core) - -if(NOT BUILD_ZVEC_CORE_ONLY) - cc_directory(db) - cc_directory(binding) -endif() +cc_directory(db) +cc_directory(binding) # ============================================================================= # Build ALL-IN-ONE C++ Shared Libraries @@ -147,15 +144,20 @@ function(zvec_add_all_in_one_shared TARGET_NAME OUTPUT_NAME) ) endfunction() -if(BUILD_ZVEC_CORE_SHARED AND NOT IOS) +if(BUILD_ZVEC_AILEGO_SHARED) + zvec_add_all_in_one_shared(zvec_ailego_shared zvec_ailego + LIBS + zvec_ailego + ) +endif() + +if(BUILD_ZVEC_CORE_SHARED) zvec_add_all_in_one_shared(zvec_core_shared zvec_core LIBS zvec_core zvec_ailego zvec_turbo ) -elseif(BUILD_ZVEC_CORE_SHARED AND IOS) - message(STATUS "Skipping zvec_core_shared on iOS") endif() if(BUILD_ZVEC_SHARED) diff --git a/src/ailego/CMakeLists.txt b/src/ailego/CMakeLists.txt index 628027c..29cf22c 100644 --- a/src/ailego/CMakeLists.txt +++ b/src/ailego/CMakeLists.txt @@ -13,21 +13,11 @@ git_version(GIT_SRCS_VER ${CMAKE_CURRENT_SOURCE_DIR}) file(GLOB_RECURSE ALL_SRCS *.cc *.c *.h) set(EXTRA_LIBS ${CMAKE_THREAD_LIBS_INIT} ${CMAKE_DL_LIBS}) -set(EXTRA_DEFS) if(UNIX AND NOT APPLE) list(APPEND EXTRA_LIBS ${LIB_RT}) endif() -if(BUILD_ZVEC_CORE_ONLY) - list(FILTER ALL_SRCS EXCLUDE REGEX ".*/buffer/parquet_hash_table\\.cc$") - list(APPEND EXTRA_DEFS ZVEC_CORE_ONLY=1) -else() - list(APPEND EXTRA_LIBS - Arrow::arrow_static - Arrow::parquet_static) -endif() - if(NOT ANDROID AND AUTO_DETECT_ARCH) if(HOST_ARCH MATCHES "^(x86|x64)$") setup_compiler_march_for_x86(MATH_MARCH_FLAG_SSE MATH_MARCH_FLAG_AVX2 MATH_MARCH_FLAG_AVX512 MATH_MARCH_FLAG_AVX512FP16) @@ -131,6 +121,5 @@ cc_library( NAME zvec_ailego STATIC STRICT PACKED SRCS ${ALL_SRCS} LIBS ${EXTRA_LIBS} - DEFS ${EXTRA_DEFS} VERSION "${GIT_SRCS_VER}" ) diff --git a/src/ailego/buffer/block_eviction_queue.cc b/src/ailego/buffer/block_eviction_queue.cc index d51ae35..eff9301 100644 --- a/src/ailego/buffer/block_eviction_queue.cc +++ b/src/ailego/buffer/block_eviction_queue.cc @@ -12,12 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include -#include - -#ifndef ZVEC_CORE_ONLY -#include -#endif +#include +#include namespace zvec { namespace ailego { @@ -42,14 +38,12 @@ bool BlockEvictionQueue::evict_single_block(BlockType &item) { } bool BlockEvictionQueue::is_valid_and_alive(const BlockType &item) { - std::shared_lock lock(valid_page_tables_mutex_); - if (valid_page_tables_.find(item.page_table) == valid_page_tables_.end()) { + std::shared_lock lock(valid_owners_mutex_); + if (item.owner == nullptr || + valid_owners_.find(item.owner) == valid_owners_.end()) { return false; } - // is_dead_block accesses entries_ under the same shared lock, so the - // VectorPageTable destructor (which holds the unique lock via set_invalid) - // cannot free entries_ while this check is in progress. - return !item.page_table->is_dead_block(item); + return !item.owner->is_dead_block(item.owner_key, item.version); } bool BlockEvictionQueue::evict_block(BlockType &item) { @@ -59,17 +53,6 @@ bool BlockEvictionQueue::evict_block(BlockType &item) { if (!ok) { return false; } - if (item.page_table == nullptr) { -#ifndef ZVEC_CORE_ONLY - if (!ParquetBufferPool::get_instance().is_dead_node(item)) { - break; - } else { - continue; - } -#else - continue; -#endif - } } while (!is_valid_and_alive(item)); return ok; } @@ -77,16 +60,10 @@ bool BlockEvictionQueue::evict_block(BlockType &item) { void BlockEvictionQueue::recycle() { BlockType item; while (MemoryLimitPool::get_instance().is_full() && evict_block(item)) { - if (item.page_table) { - std::shared_lock lock(valid_page_tables_mutex_); - if (valid_page_tables_.find(item.page_table) != - valid_page_tables_.end()) { - item.page_table->evict_block(item.vector_block.first); - } -#ifndef ZVEC_CORE_ONLY - } else { - ParquetBufferPool::get_instance().evict(item.parquet_buffer_block.first); -#endif + std::shared_lock lock(valid_owners_mutex_); + if (item.owner != nullptr && + valid_owners_.find(item.owner) != valid_owners_.end()) { + item.owner->evict_block(item.owner_key); } } } @@ -127,7 +104,7 @@ bool MemoryLimitPool::try_acquire_buffer(const size_t buffer_size, return true; } -void MemoryLimitPool::acquire_parquet(const size_t buffer_size) { +void MemoryLimitPool::charge_external(const size_t buffer_size) { size_t expected, desired; do { expected = used_size_.load(); @@ -145,7 +122,7 @@ void MemoryLimitPool::release_buffer(char *buffer, const size_t buffer_size) { ailego_free(buffer); } -void MemoryLimitPool::release_parquet(const size_t buffer_size) { +void MemoryLimitPool::release_external(const size_t buffer_size) { size_t expected, desired; do { expected = used_size_.load(); diff --git a/src/ailego/buffer/parquet_hash_table.cc b/src/ailego/buffer/parquet_hash_table.cc deleted file mode 100644 index ab51984..0000000 --- a/src/ailego/buffer/parquet_hash_table.cc +++ /dev/null @@ -1,250 +0,0 @@ -// Copyright 2025-present the zvec project -// -// Licensed 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 -#include -#include -#include -#include -#include -#include - -namespace zvec { -namespace ailego { - -ParquetBufferID::ParquetBufferID(std::string &filename, int column, - int row_group) - : filename(filename), column(column), row_group(row_group) { - struct stat file_stat; - if (stat(filename.c_str(), &file_stat) == 0) { - // file_stat.st_ino contains the inode number - // file_stat.st_dev contains the device ID - // Together they uniquely identify a file - file_id = file_stat.st_ino; - std::filesystem::path p(filename); - auto ftime = std::filesystem::last_write_time(p); - mtime = static_cast(ftime.time_since_epoch().count()); - } -} - -const std::string ParquetBufferID::to_string() const { - std::string msg{"Buffer["}; - msg += "parquet: " + filename + "[" + std::to_string(file_id) + "]" + - ", column: " + std::to_string(column) + - ", row_group: " + std::to_string(row_group); - msg += ", mtime: " + std::to_string(mtime); - msg += "]"; - return msg; -} - -ParquetBufferContextHandle::ParquetBufferContextHandle( - const ParquetBufferContextHandle &handle_) - : buffer_id_(handle_.buffer_id_), arrow_(handle_.arrow_) { - if (arrow_) { - ParquetBufferPool::get_instance().acquire_locked(buffer_id_); - } -} - -ParquetBufferContextHandle::~ParquetBufferContextHandle() { - if (arrow_) { - ParquetBufferPool::get_instance().release(buffer_id_); - } -} - -arrow::Status ParquetBufferPool::acquire(ParquetBufferID buffer_id, - ParquetBufferContext &context) { - // TODO: file handler and memory pool can be optimized - arrow::MemoryPool *mem_pool = arrow::default_memory_pool(); - - // Open file - std::shared_ptr input; - const auto &file_name = buffer_id.filename; - ARROW_ASSIGN_OR_RAISE(input, arrow::io::ReadableFile::Open(file_name)); - - // Open reader - std::unique_ptr reader; - ARROW_ASSIGN_OR_RAISE(reader, parquet::arrow::OpenFile(input, mem_pool)); - - // Perform read - int row_group = buffer_id.row_group; - int column = buffer_id.column; - auto s = reader->RowGroup(row_group)->Column(column)->Read(&context.arrow); - if (!s.ok()) { - LOG_ERROR("Failed to read parquet file[%s]", file_name.c_str()); - context.arrow = nullptr; - return s; - } - - context.size = 0; - context.arrow_refs.clear(); - // Compute the memory usage and hijack Arrow's buffers with our - // implementation - for (auto &array : context.arrow->chunks()) { - auto &buffers = array->data()->buffers; - for (size_t buf_idx = 0; buf_idx < buffers.size(); ++buf_idx) { - if (buffers[buf_idx] == nullptr) { - continue; - } - // Keep references to original buffers to prevent premature deletion - context.arrow_refs.emplace_back(buffers[buf_idx]); - context.size += buffers[buf_idx]->capacity(); - // Create hijacked buffer with custom deleter that notifies us when - // Arrow is finished with the buffer - std::shared_ptr hijacked_buffer( - buffers[buf_idx].get(), ArrowBufferDeleter(this, buffer_id)); - buffers[buf_idx] = hijacked_buffer; - } - } - - return arrow::Status::OK(); -} - -ParquetBufferContextHandle ParquetBufferPool::acquire_buffer( - ParquetBufferID buffer_id) { - std::shared_ptr arrow{nullptr}; - { - std::shared_lock lock(table_mutex_); - auto iter = table_.find(buffer_id); - if (iter != table_.end()) { - arrow = acquire(buffer_id); - if (arrow != nullptr) { - return ParquetBufferContextHandle(buffer_id, arrow); - } - } - } - { - bool found = !MemoryLimitPool::get_instance().is_full(); - if (!found) { - for (int i = 0; i < 5; i++) { - BlockEvictionQueue::get_instance().recycle(); - found = !MemoryLimitPool::get_instance().is_full(); - if (found) { - break; - } - } - } - if (!found) { - LOG_ERROR("Failed to acquire parquet buffer: %s", - buffer_id.to_string().c_str()); - return ParquetBufferContextHandle(); - } - std::unique_lock lock(table_mutex_); - if (acquire(buffer_id, table_[buffer_id]).ok()) { - MemoryLimitPool::get_instance().acquire_parquet(table_[buffer_id].size); - arrow = set_block_acquired(buffer_id); - return ParquetBufferContextHandle(buffer_id, arrow); - } else { - LOG_ERROR("Failed to acquire parquet buffer: %s", - buffer_id.to_string().c_str()); - return ParquetBufferContextHandle(); - } - } -} - -std::shared_ptr ParquetBufferPool::set_block_acquired( - ParquetBufferID buffer_id) { - ParquetBufferContext &context = table_[buffer_id]; - while (true) { - int current_count = context.ref_count.load(std::memory_order_relaxed); - if (current_count >= 0) { - if (context.ref_count.compare_exchange_weak( - current_count, current_count + 1, std::memory_order_acq_rel, - std::memory_order_acquire)) { - return context.arrow; - } - } else { - if (context.ref_count.compare_exchange_weak(current_count, 1, - std::memory_order_acq_rel, - std::memory_order_acquire)) { - context.load_count.fetch_add(1, std::memory_order_relaxed); - return context.arrow; - } - } - } -} - -std::shared_ptr ParquetBufferPool::acquire( - ParquetBufferID buffer_id) { - auto iter = table_.find(buffer_id); - if (iter == table_.end()) { - return nullptr; - } - ParquetBufferContext &context = table_[buffer_id]; - while (true) { - int current_count = context.ref_count.load(std::memory_order_acquire); - if (current_count < 0) { - return nullptr; - } - if (context.ref_count.compare_exchange_weak( - current_count, current_count + 1, std::memory_order_acq_rel, - std::memory_order_acquire)) { - if (current_count == 0) { - context.load_count.fetch_add(1, std::memory_order_relaxed); - } - return context.arrow; - } - } - return nullptr; -} - -std::shared_ptr ParquetBufferPool::acquire_locked( - ParquetBufferID buffer_id) { - std::shared_lock lock(table_mutex_); - return acquire(buffer_id); -} - -void ParquetBufferPool::release(ParquetBufferID buffer_id) { - std::shared_lock lock(table_mutex_); - auto iter = table_.find(buffer_id); - if (iter == table_.end()) { - return; - } - ParquetBufferContext &context = table_[buffer_id]; - if (context.ref_count.fetch_sub(1, std::memory_order_release) == 1) { - std::atomic_thread_fence(std::memory_order_acquire); - BlockEvictionQueue::BlockType block; - block.parquet_buffer_block.first = buffer_id; - block.parquet_buffer_block.second = context.load_count.load(); - BlockEvictionQueue::get_instance().add_single_block(block, 0); - } -} - -void ParquetBufferPool::evict(ParquetBufferID buffer_id) { - std::unique_lock lock(table_mutex_); - auto iter = table_.find(buffer_id); - if (iter == table_.end()) { - return; - } - ParquetBufferContext &context = table_[buffer_id]; - int expected = 0; - if (context.ref_count.compare_exchange_strong( - expected, std::numeric_limits::min())) { - MemoryLimitPool::get_instance().release_parquet(context.size); - context.arrow = nullptr; - context.arrow_refs.clear(); - } -} - -bool ParquetBufferPool::is_dead_node(BlockEvictionQueue::BlockType &block) { - std::shared_lock lock(table_mutex_); - auto iter = table_.find(block.parquet_buffer_block.first); - if (iter == table_.end()) { - return true; - } - return iter->second.load_count.load() != block.parquet_buffer_block.second; -} - -} // namespace ailego -} // namespace zvec \ No newline at end of file diff --git a/src/ailego/buffer/vector_page_table.cc b/src/ailego/buffer/vector_page_table.cc index c9296d6..3318db1 100644 --- a/src/ailego/buffer/vector_page_table.cc +++ b/src/ailego/buffer/vector_page_table.cc @@ -18,7 +18,7 @@ #include #include #include -#include +#include #if defined(_MSC_VER) #ifndef NOMINMAX @@ -167,9 +167,9 @@ void VectorPageTable::release_block(block_id_t block_id) { std::memory_order_acq_rel, std::memory_order_relaxed)) { BlockEvictionQueue::BlockType block; - block.page_table = this; - block.vector_block.first = block_id; - block.vector_block.second = 0; + block.owner = this; + block.owner_key = block_id; + block.version = 0; BlockEvictionQueue::get_instance().add_single_block(block, 0); } } @@ -618,4 +618,4 @@ void VecBufferPoolHandle::acquire_one(block_id_t block_id) { } } // namespace ailego -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_streamer_entity.h b/src/core/algorithm/hnsw/hnsw_streamer_entity.h index 013ad8a..19f8ba1 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer_entity.h +++ b/src/core/algorithm/hnsw/hnsw_streamer_entity.h @@ -1099,4 +1099,4 @@ class HnswContiguousStreamerEntity : public HnswMmapStreamerEntity { }; } // namespace core -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/db/CMakeLists.txt b/src/db/CMakeLists.txt index 7808b29..5dfb081 100644 --- a/src/db/CMakeLists.txt +++ b/src/db/CMakeLists.txt @@ -46,6 +46,7 @@ cc_library( FastPFOR cppjieba Arrow::arrow_static + Arrow::parquet_static Arrow::arrow_compute Arrow::arrow_dataset Arrow::arrow_acero diff --git a/src/db/common/global_resource.cc b/src/db/common/global_resource.cc index d6c7a4b..711bf3d 100644 --- a/src/db/common/global_resource.cc +++ b/src/db/common/global_resource.cc @@ -30,4 +30,4 @@ void GlobalResource::initialize() { }); } -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/db/index/CMakeLists.txt b/src/db/index/CMakeLists.txt index d4efc32..8231718 100644 --- a/src/db/index/CMakeLists.txt +++ b/src/db/index/CMakeLists.txt @@ -25,6 +25,7 @@ cc_library( rocksdb core_interface Arrow::arrow_static + Arrow::parquet_static Arrow::arrow_compute Arrow::arrow_dataset cppjieba diff --git a/src/db/index/storage/bufferpool_forward_store.cc b/src/db/index/storage/bufferpool_forward_store.cc index 36aefe4..2bd0bd6 100644 --- a/src/db/index/storage/bufferpool_forward_store.cc +++ b/src/db/index/storage/bufferpool_forward_store.cc @@ -21,10 +21,10 @@ #include #include #include -#include #include #include "db/index/storage/store_helper.h" #include "lazy_record_batch_reader.h" +#include "parquet_buffer_pool.h" namespace zvec { @@ -191,9 +191,9 @@ TablePtr BufferPoolForwardStore::fetch(const std::vector &columns, for (const auto &[rg_id, pairs] : rg_to_local) { for (size_t i = 0; i < col_indices.size(); ++i) { int col_idx = col_indices[i]; - auto buffer_id = ailego::ParquetBufferID(file_path_, col_idx, rg_id); + auto buffer_id = ParquetBufferID(file_path_, col_idx, rg_id); auto buffer_handle = - ailego::ParquetBufferPool::get_instance().acquire_buffer(buffer_id); + ParquetBufferPool::get_instance().acquire_buffer(buffer_id); std::shared_ptr col_chunked_array = buffer_handle.data(); if (!col_chunked_array) { @@ -317,9 +317,9 @@ ExecBatchPtr BufferPoolForwardStore::fetch( std::vector scalars; for (size_t i = 0; i < col_indices.size(); ++i) { int col_idx = col_indices[i]; - auto buffer_id = ailego::ParquetBufferID(file_path_, col_idx, rg_id); + auto buffer_id = ParquetBufferID(file_path_, col_idx, rg_id); auto buffer_handle = - ailego::ParquetBufferPool::get_instance().acquire_buffer(buffer_id); + ParquetBufferPool::get_instance().acquire_buffer(buffer_id); std::shared_ptr col_chunked_array = buffer_handle.data(); @@ -379,4 +379,4 @@ RecordBatchReaderPtr BufferPoolForwardStore::scan( physic_schema_, file_path_); } -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/db/index/storage/bufferpool_forward_store.h b/src/db/index/storage/bufferpool_forward_store.h index 45ddc3a..672ac69 100644 --- a/src/db/index/storage/bufferpool_forward_store.h +++ b/src/db/index/storage/bufferpool_forward_store.h @@ -124,4 +124,4 @@ class BufferPoolForwardStore std::vector row_group_row_nums_; }; -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/db/index/storage/lazy_record_batch_reader.h b/src/db/index/storage/lazy_record_batch_reader.h index 8102995..2edcb89 100644 --- a/src/db/index/storage/lazy_record_batch_reader.h +++ b/src/db/index/storage/lazy_record_batch_reader.h @@ -16,8 +16,8 @@ #include #include -#include #include "db/common/constants.h" +#include "parquet_buffer_pool.h" namespace zvec { @@ -128,9 +128,9 @@ class ParquetRecordBatchReader : public arrow::RecordBatchReader { if (with_cache_) { for (size_t col_idx = 0; col_idx < col_indices_.size(); ++col_idx) { auto buffer_id = - ailego::ParquetBufferID(file_path_, col_indices_[col_idx], rg_id); + ParquetBufferID(file_path_, col_indices_[col_idx], rg_id); auto buffer_handle = - ailego::ParquetBufferPool::get_instance().acquire_buffer(buffer_id); + ParquetBufferPool::get_instance().acquire_buffer(buffer_id); std::shared_ptr col_chunked_array = buffer_handle.data(); if (col_chunked_array) { diff --git a/src/db/index/storage/parquet_buffer_pool.cc b/src/db/index/storage/parquet_buffer_pool.cc new file mode 100644 index 0000000..03e454a --- /dev/null +++ b/src/db/index/storage/parquet_buffer_pool.cc @@ -0,0 +1,142 @@ +// Copyright 2025-present the zvec project +// +// Licensed 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 "parquet_buffer_pool.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace zvec { + +ParquetBufferID::ParquetBufferID(const std::string &filename, int column, + int row_group) + : filename(filename), column(column), row_group(row_group) { + struct stat file_stat; + if (stat(filename.c_str(), &file_stat) == 0) { + file_id = file_stat.st_ino; + std::filesystem::path p(filename); + auto ftime = std::filesystem::last_write_time(p); + mtime = static_cast(ftime.time_since_epoch().count()); + } +} + +const std::string ParquetBufferID::to_string() const { + std::string msg{"Buffer["}; + msg += "parquet: " + filename + "[" + std::to_string(file_id) + "]" + + ", column: " + std::to_string(column) + + ", row_group: " + std::to_string(row_group); + msg += ", mtime: " + std::to_string(mtime); + msg += "]"; + return msg; +} + +ParquetBufferContextHandle::ParquetBufferContextHandle( + const ParquetBufferContextHandle &handle) + : buffer_id_(handle.buffer_id_) { + if (handle.arrow_) { + arrow_ = ParquetBufferPool::get_instance().retain(buffer_id_); + } +} + +ParquetBufferContextHandle::~ParquetBufferContextHandle() { + if (arrow_) { + ParquetBufferPool::get_instance().release(buffer_id_); + } +} + +bool detail::ParquetBufferLoader::load(const ParquetBufferID &buffer_id, + ParquetBufferPayload &payload, + size_t &size) { + arrow::MemoryPool *mem_pool = arrow::default_memory_pool(); + + std::shared_ptr input; + const auto &file_name = buffer_id.filename; + auto input_result = arrow::io::ReadableFile::Open(file_name); + if (!input_result.ok()) { + LOG_ERROR("Failed to open parquet file[%s]: %s", file_name.c_str(), + input_result.status().ToString().c_str()); + return false; + } + input = *input_result; + + std::unique_ptr reader; + auto reader_result = parquet::arrow::OpenFile(input, mem_pool); + if (!reader_result.ok()) { + LOG_ERROR("Failed to create parquet reader[%s]: %s", file_name.c_str(), + reader_result.status().ToString().c_str()); + return false; + } + reader = std::move(*reader_result); + + int row_group = buffer_id.row_group; + int column = buffer_id.column; + auto s = reader->RowGroup(row_group)->Column(column)->Read(&payload.arrow); + if (!s.ok()) { + LOG_ERROR("Failed to read parquet file[%s]: %s", file_name.c_str(), + s.ToString().c_str()); + payload.arrow = nullptr; + return false; + } + + size = 0; + payload.arrow_refs.clear(); + for (auto &array : payload.arrow->chunks()) { + auto &buffers = array->data()->buffers; + for (size_t buf_idx = 0; buf_idx < buffers.size(); ++buf_idx) { + if (buffers[buf_idx] == nullptr) { + continue; + } + payload.arrow_refs.emplace_back(buffers[buf_idx]); + size += buffers[buf_idx]->capacity(); + std::shared_ptr hijacked_buffer(buffers[buf_idx].get(), + [](arrow::Buffer *) {}); + buffers[buf_idx] = hijacked_buffer; + } + } + + return true; +} + +void detail::ParquetBufferLoader::clear(ParquetBufferPayload &payload) const { + payload.arrow = nullptr; + payload.arrow_refs.clear(); +} + +ParquetBufferContextHandle ParquetBufferPool::acquire_buffer( + ParquetBufferID buffer_id) { + auto arrow = cache_.acquire(buffer_id); + if (!arrow) { + LOG_ERROR("Failed to acquire parquet buffer: %s", + buffer_id.to_string().c_str()); + return ParquetBufferContextHandle(); + } + return ParquetBufferContextHandle(buffer_id, arrow); +} + +void ParquetBufferPool::release(ParquetBufferID buffer_id) { + cache_.release(buffer_id); +} + +std::shared_ptr ParquetBufferPool::retain( + ParquetBufferID buffer_id) { + return cache_.retain(buffer_id); +} + +} // namespace zvec diff --git a/src/include/zvec/ailego/buffer/parquet_hash_table.h b/src/db/index/storage/parquet_buffer_pool.h similarity index 56% rename from src/include/zvec/ailego/buffer/parquet_hash_table.h rename to src/db/index/storage/parquet_buffer_pool.h index 4db1a8f..f04871c 100644 --- a/src/include/zvec/ailego/buffer/parquet_hash_table.h +++ b/src/db/index/storage/parquet_buffer_pool.h @@ -12,40 +12,37 @@ // See the License for the specific language governing permissions and // limitations under the License. - #pragma once #include -#include -#include #include #include #include -#include -#include +#include #include -#include -#include -#include -#include "block_eviction_queue.h" +#include namespace arrow { -class ChunkedArray; -class Array; -class DataType; -class Scalar; -template -class Result; -class Status; class Buffer; +class ChunkedArray; } // namespace arrow namespace zvec { -namespace ailego { -class BlockEvictionQueue; +struct ParquetBufferID { + std::string filename; + int column{0}; + int row_group{0}; + uint64_t file_id{0}; + long mtime{0}; -struct IDHash { + ParquetBufferID() = default; + ParquetBufferID(const std::string &filename, int column, int row_group); + + const std::string to_string() const; +}; + +struct ParquetBufferIDHash { size_t operator()(const ParquetBufferID &buffer_id) const { size_t hash = std::hash{}(1); hash = hash ^ (std::hash{}(buffer_id.file_id)); @@ -55,7 +52,7 @@ struct IDHash { } }; -struct IDEqual { +struct ParquetBufferIDEqual { bool operator()(const ParquetBufferID &a, const ParquetBufferID &b) const { if (a.filename != b.filename) { return false; @@ -70,28 +67,39 @@ struct IDEqual { } }; -struct ParquetBufferContext { - // A shared pointer to the buffers allocated for arrow parquet data +namespace detail { + +struct ParquetBufferPayload { std::shared_ptr arrow{nullptr}; - - // Guard original arrow buffers to prevent premature deletion std::vector> arrow_refs{}; - - size_t size; - alignas(64) std::atomic ref_count{std::numeric_limits::min()}; - alignas(64) std::atomic load_count{0}; }; +class ParquetBufferLoader { + public: + using Value = std::shared_ptr; + + bool load(const ParquetBufferID &buffer_id, ParquetBufferPayload &payload, + size_t &size); + + Value value(const ParquetBufferPayload &payload) const { + return payload.arrow; + } + + void clear(ParquetBufferPayload &payload) const; +}; + +} // namespace detail + class ParquetBufferContextHandle { public: - ParquetBufferContextHandle() {} - ParquetBufferContextHandle(ParquetBufferID &buffer_id, + ParquetBufferContextHandle() = default; + ParquetBufferContextHandle(const ParquetBufferID &buffer_id, std::shared_ptr arrow) - : buffer_id_(buffer_id), arrow_(arrow) {} - ParquetBufferContextHandle(const ParquetBufferContextHandle &handle_); - ParquetBufferContextHandle(ParquetBufferContextHandle &&handle_) - : buffer_id_(std::move(handle_.buffer_id_)), - arrow_(std::move(handle_.arrow_)) {} + : buffer_id_(buffer_id), arrow_(std::move(arrow)) {} + ParquetBufferContextHandle(const ParquetBufferContextHandle &handle); + ParquetBufferContextHandle(ParquetBufferContextHandle &&handle) + : buffer_id_(std::move(handle.buffer_id_)), + arrow_(std::move(handle.arrow_)) {} ~ParquetBufferContextHandle(); @@ -108,40 +116,8 @@ class ParquetBufferPool { public: typedef std::shared_ptr Pointer; - struct ArrowBufferDeleter { - explicit ArrowBufferDeleter(ParquetBufferPool *c, ParquetBufferID i) - : pool(c), id(i) {} - ParquetBufferPool *pool; - ParquetBufferID id; - // Only reduces the reference count but does not actually release the - // buffer, since the buffer memory is managed by the ParquetBufferPool. - void operator()(arrow::Buffer *) { - return; - } - }; - - using Table = std::unordered_map; - - arrow::Status acquire(ParquetBufferID buffer_id, - ParquetBufferContext &context); - ParquetBufferContextHandle acquire_buffer(ParquetBufferID buffer_id); - std::shared_ptr set_block_acquired( - ParquetBufferID buffer_id); - - std::shared_ptr acquire(ParquetBufferID buffer_id); - - std::shared_ptr acquire_locked( - ParquetBufferID buffer_id); - - void release(ParquetBufferID buffer_id); - - void evict(ParquetBufferID buffer_id); - - bool is_dead_node(BlockEvictionQueue::BlockType &block); - static ParquetBufferPool &get_instance() { static ParquetBufferPool instance; return instance; @@ -153,12 +129,21 @@ class ParquetBufferPool { ParquetBufferPool &operator=(ParquetBufferPool &&) = delete; private: + friend class ParquetBufferContextHandle; + + using Cache = + ailego::ExternalCache; + ParquetBufferPool() = default; + std::shared_ptr retain(ParquetBufferID buffer_id); + + void release(ParquetBufferID buffer_id); + private: - Table table_; - std::shared_mutex table_mutex_; + Cache cache_; }; -} // namespace ailego -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/include/zvec/ailego/buffer/block_eviction_queue.h b/src/include/zvec/ailego/buffer/block_eviction_queue.h index 2c78b99..fa5aff2 100644 --- a/src/include/zvec/ailego/buffer/block_eviction_queue.h +++ b/src/include/zvec/ailego/buffer/block_eviction_queue.h @@ -22,7 +22,6 @@ #include #include #include -#include #include #include #include @@ -34,7 +33,6 @@ #include #include #include -#include #include "concurrentqueue.h" #if defined(_MSC_VER) @@ -44,28 +42,25 @@ namespace zvec { namespace ailego { -class VectorPageTable; - +using eviction_key_t = size_t; using block_id_t = size_t; using version_t = size_t; -struct ParquetBufferID { - std::string filename; - int column; - int row_group; - uint64_t file_id; - long mtime; - ParquetBufferID() = default; - ParquetBufferID(std::string &filename, int column, int row_group); - const std::string to_string() const; +class EvictableBlockOwner { + public: + virtual ~EvictableBlockOwner() = default; + + virtual bool is_dead_block(eviction_key_t owner_key, version_t version) = 0; + + virtual void evict_block(eviction_key_t owner_key) = 0; }; class BlockEvictionQueue { public: struct BlockType { - std::pair vector_block; - std::pair parquet_buffer_block; - VectorPageTable *page_table{nullptr}; + eviction_key_t owner_key{0}; + version_t version{0}; + EvictableBlockOwner *owner{nullptr}; }; typedef moodycamel::ConcurrentQueue ConcurrentQueue; @@ -88,24 +83,24 @@ class BlockEvictionQueue { // void clear_dead_node(); - bool is_valid(VectorPageTable *page_table) { - std::shared_lock lock(valid_page_tables_mutex_); - return valid_page_tables_.find(page_table) != valid_page_tables_.end(); + bool is_valid(EvictableBlockOwner *owner) { + std::shared_lock lock(valid_owners_mutex_); + return valid_owners_.find(owner) != valid_owners_.end(); } - void set_valid(VectorPageTable *page_table) { - std::unique_lock lock(valid_page_tables_mutex_); - valid_page_tables_.insert(page_table); + void set_valid(EvictableBlockOwner *owner) { + std::unique_lock lock(valid_owners_mutex_); + valid_owners_.insert(owner); } - void set_invalid(VectorPageTable *page_table) { - std::unique_lock lock(valid_page_tables_mutex_); - valid_page_tables_.erase(page_table); + void set_invalid(EvictableBlockOwner *owner) { + std::unique_lock lock(valid_owners_mutex_); + valid_owners_.erase(owner); } - // Atomically checks under the shared lock that the page table is still valid - // AND the block version has not been superseded, preventing TOCTOU races - // when a VectorPageTable is concurrently destroyed. + // Atomically checks under the shared lock that the owner is still valid AND + // the block version has not been superseded, preventing TOCTOU races when an + // owner is concurrently destroyed. bool is_valid_and_alive(const BlockType &item); void recycle(); @@ -119,8 +114,8 @@ class BlockEvictionQueue { constexpr static size_t CACHE_QUEUE_NUM = 3; size_t evict_batch_size_{0}; std::vector evict_queues_; - std::unordered_set valid_page_tables_; - std::shared_mutex valid_page_tables_mutex_; + std::unordered_set valid_owners_; + std::shared_mutex valid_owners_mutex_; }; class MemoryLimitPool { @@ -138,11 +133,11 @@ class MemoryLimitPool { bool try_acquire_buffer(const size_t buffer_size, char *&buffer); - void acquire_parquet(const size_t buffer_size); + void charge_external(const size_t buffer_size); void release_buffer(char *buffer, const size_t buffer_size); - void release_parquet(const size_t buffer_size); + void release_external(const size_t buffer_size); bool is_full(); @@ -155,4 +150,4 @@ class MemoryLimitPool { }; } // namespace ailego -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/include/zvec/ailego/buffer/external_cache.h b/src/include/zvec/ailego/buffer/external_cache.h new file mode 100644 index 0000000..671ef45 --- /dev/null +++ b/src/include/zvec/ailego/buffer/external_cache.h @@ -0,0 +1,213 @@ +// Copyright 2025-present the zvec project +// +// Licensed 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 +#include +#include + +namespace zvec { +namespace ailego { + +template +class ExternalCache : public EvictableBlockOwner { + public: + using Value = typename Loader::Value; + + ExternalCache() { + BlockEvictionQueue::get_instance().set_valid(this); + } + + explicit ExternalCache(Loader loader) : loader_(std::move(loader)) { + BlockEvictionQueue::get_instance().set_valid(this); + } + + ~ExternalCache() { + BlockEvictionQueue::get_instance().set_invalid(this); + } + + ExternalCache(const ExternalCache &) = delete; + ExternalCache &operator=(const ExternalCache &) = delete; + ExternalCache(ExternalCache &&) = delete; + ExternalCache &operator=(ExternalCache &&) = delete; + + Value acquire(const Key &key) { + { + std::shared_lock lock(mutex_); + auto iter = table_.find(key); + if (iter != table_.end()) { + Value value = acquire_loaded(iter->second); + if (value) { + return value; + } + } + } + + if (!ensure_capacity()) { + return Value{}; + } + + std::unique_lock lock(mutex_); + auto iter = table_.find(key); + if (iter != table_.end()) { + Value value = acquire_loaded(iter->second); + if (value) { + return value; + } + } else { + auto inserted = table_.try_emplace(key); + iter = inserted.first; + iter->second.owner_key = next_owner_key_++; + owner_keys_.emplace(iter->second.owner_key, key); + } + + Entry &entry = iter->second; + size_t size = 0; + if (!loader_.load(key, entry.payload, size)) { + return Value{}; + } + + entry.size = size; + MemoryLimitPool::get_instance().charge_external(entry.size); + entry.generation.fetch_add(1, std::memory_order_relaxed); + entry.ref_count.store(1, std::memory_order_release); + return loader_.value(entry.payload); + } + + Value retain(const Key &key) { + std::shared_lock lock(mutex_); + auto iter = table_.find(key); + if (iter == table_.end()) { + return Value{}; + } + return acquire_loaded(iter->second); + } + + void release(const Key &key) { + std::shared_lock lock(mutex_); + auto iter = table_.find(key); + if (iter == table_.end()) { + return; + } + + Entry &entry = iter->second; + if (entry.ref_count.fetch_sub(1, std::memory_order_release) == 1) { + std::atomic_thread_fence(std::memory_order_acquire); + BlockEvictionQueue::BlockType block; + block.owner = this; + block.owner_key = entry.owner_key; + block.version = entry.generation.load(std::memory_order_relaxed); + BlockEvictionQueue::get_instance().add_single_block(block, 0); + } + } + + bool is_dead_block(eviction_key_t owner_key, version_t version) override { + std::shared_lock lock(mutex_); + auto key_iter = owner_keys_.find(owner_key); + if (key_iter == owner_keys_.end()) { + return true; + } + + auto iter = table_.find(key_iter->second); + if (iter == table_.end()) { + return true; + } + return iter->second.generation.load(std::memory_order_relaxed) != version; + } + + void evict_block(eviction_key_t owner_key) override { + std::unique_lock lock(mutex_); + auto key_iter = owner_keys_.find(owner_key); + if (key_iter == owner_keys_.end()) { + return; + } + + auto iter = table_.find(key_iter->second); + if (iter == table_.end()) { + return; + } + + Entry &entry = iter->second; + int expected = 0; + if (entry.ref_count.compare_exchange_strong( + expected, std::numeric_limits::min())) { + MemoryLimitPool::get_instance().release_external(entry.size); + entry.size = 0; + loader_.clear(entry.payload); + } + } + + private: + struct Entry { + Payload payload{}; + size_t size{0}; + eviction_key_t owner_key{0}; + alignas(64) std::atomic ref_count{std::numeric_limits::min()}; + alignas(64) std::atomic generation{0}; + }; + + Value acquire_loaded(Entry &entry) { + while (true) { + int current_count = entry.ref_count.load(std::memory_order_acquire); + if (current_count < 0) { + return Value{}; + } + if (entry.ref_count.compare_exchange_weak( + current_count, current_count + 1, std::memory_order_acq_rel, + std::memory_order_acquire)) { + if (current_count == 0) { + entry.generation.fetch_add(1, std::memory_order_relaxed); + } + return loader_.value(entry.payload); + } + } + } + + bool ensure_capacity() { + bool found = !MemoryLimitPool::get_instance().is_full(); + if (found) { + return true; + } + + for (int i = 0; i < kRecycleAttempts; ++i) { + BlockEvictionQueue::get_instance().recycle(); + found = !MemoryLimitPool::get_instance().is_full(); + if (found) { + return true; + } + } + return false; + } + + private: + static constexpr int kRecycleAttempts = 5; + + using Table = std::unordered_map; + + Loader loader_{}; + Table table_; + std::unordered_map owner_keys_; + eviction_key_t next_owner_key_{1}; + std::shared_mutex mutex_; +}; + +} // namespace ailego +} // namespace zvec diff --git a/src/include/zvec/ailego/buffer/vector_page_table.h b/src/include/zvec/ailego/buffer/vector_page_table.h index 8bcc13e..02d19bb 100644 --- a/src/include/zvec/ailego/buffer/vector_page_table.h +++ b/src/include/zvec/ailego/buffer/vector_page_table.h @@ -45,7 +45,7 @@ namespace ailego { extern const size_t kVectorPageSize; -class VectorPageTable { +class VectorPageTable : public EvictableBlockOwner { struct Entry { std::atomic ref_count; std::atomic in_evict_queue; @@ -94,7 +94,7 @@ class VectorPageTable { void release_block(block_id_t block_id); - void evict_block(block_id_t block_id); + void evict_block(block_id_t block_id) override; char *set_block_acquired(block_id_t block_id, char *buffer, size_t file_offset); @@ -145,8 +145,9 @@ class VectorPageTable { return entry_at(block_id).ref_count.load(std::memory_order_relaxed) <= 0; } - inline bool is_dead_block(BlockEvictionQueue::BlockType block) const { - const Entry &e = entry_at(block.vector_block.first); + inline bool is_dead_block(block_id_t block_id, + version_t /*version*/) override { + const Entry &e = entry_at(block_id); return !e.in_evict_queue.load(std::memory_order_relaxed); } @@ -297,4 +298,4 @@ class VecBufferPoolHandle { }; } // namespace ailego -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/tests/ailego/CMakeLists.txt b/tests/ailego/CMakeLists.txt index 9165643..9a52af9 100644 --- a/tests/ailego/CMakeLists.txt +++ b/tests/ailego/CMakeLists.txt @@ -9,8 +9,6 @@ foreach(CC_SRCS ${ALL_TEST_SRCS}) cc_gtest( NAME ${CC_TARGET} STRICT LIBS zvec_ailego - Arrow::arrow_depends - Arrow::parquet_static SRCS ${CC_SRCS} ) cc_test_suite(zvec_ailego ${CC_TARGET}) diff --git a/thirdparty/CMakeLists.txt b/thirdparty/CMakeLists.txt index 8c73e62..c7b2272 100644 --- a/thirdparty/CMakeLists.txt +++ b/thirdparty/CMakeLists.txt @@ -18,10 +18,6 @@ add_subdirectory(magic_enum magic_enum EXCLUDE_FROM_ALL) add_subdirectory(RaBitQ-Library RaBitQ-Library EXCLUDE_FROM_ALL) add_subdirectory(googletest googletest EXCLUDE_FROM_ALL) -if(BUILD_ZVEC_CORE_ONLY) - return() -endif() - add_subdirectory(arrow arrow EXCLUDE_FROM_ALL) add_subdirectory(gflags gflags EXCLUDE_FROM_ALL) add_subdirectory(glog glog EXCLUDE_FROM_ALL)