xiaoxue_admin
/
zvec
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: refactor parquet buffer cache ownership (#492)

This commit is contained in:
feihongxu0824 2026-06-16 10:47:07 +08:00 committed by GitHub
parent 2dbac5b348
commit 841b1ed3f9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
24 changed files with 561 additions and 634 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
.github/cmake/subproject-integration/CMakeLists.txt vendored
View File

@ -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)

58
.github/workflows/05-windows-build.yml vendored
View File

@ -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"
}
}
}

56
.github/workflows/08-cmake-subproject-integration.yml vendored
View File

@ -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

25
CMakeLists.txt
View File

@ -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)

104
examples/c++/CMakeLists.txt
View File

@ -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$<$<CONFIG:Debug>: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=<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)

18
src/CMakeLists.txt
View File

@ -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)

11
src/ailego/CMakeLists.txt
View File

@ -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}"
)

47
src/ailego/buffer/block_eviction_queue.cc
View File

@ -12,12 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include <zvec/ailego/buffer/vector_page_table.h>
#include <zvec/core/framework/index_logger.h>
#ifndef ZVEC_CORE_ONLY
#include <zvec/ailego/buffer/parquet_hash_table.h>
#endif
#include <zvec/ailego/buffer/block_eviction_queue.h>
#include <zvec/ailego/logger/logger.h>
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<std::shared_mutex> lock(valid_page_tables_mutex_);
if (valid_page_tables_.find(item.page_table) == valid_page_tables_.end()) {
std::shared_lock<std::shared_mutex> 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<std::shared_mutex> 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<std::shared_mutex> 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();

250
src/ailego/buffer/parquet_hash_table.cc
View File

@ -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 <arrow/array/array_binary.h>
#include <arrow/io/file.h>
#include <arrow/ipc/reader.h>
#include <arrow/pretty_print.h>
#include <arrow/result.h>
#include <arrow/table.h>
#include <parquet/arrow/reader.h>
#include <zvec/ailego/buffer/parquet_hash_table.h>
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<std::uint64_t>(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<arrow::io::RandomAccessFile> input;
const auto &file_name = buffer_id.filename;
ARROW_ASSIGN_OR_RAISE(input, arrow::io::ReadableFile::Open(file_name));
// Open reader
std::unique_ptr<parquet::arrow::FileReader> 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<arrow::Buffer> 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::ChunkedArray> arrow{nullptr};
{
std::shared_lock<std::shared_mutex> 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<std::shared_mutex> 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<arrow::ChunkedArray> 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<arrow::ChunkedArray> 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<arrow::ChunkedArray> ParquetBufferPool::acquire_locked(
ParquetBufferID buffer_id) {
std::shared_lock<std::shared_mutex> lock(table_mutex_);
return acquire(buffer_id);
}
void ParquetBufferPool::release(ParquetBufferID buffer_id) {
std::shared_lock<std::shared_mutex> 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<std::shared_mutex> 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<int>::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<std::shared_mutex> 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

10
src/ailego/buffer/vector_page_table.cc
View File

@ -18,7 +18,7 @@
#include <thread>
#include <ailego/utility/memory_helper.h>
#include <zvec/ailego/buffer/vector_page_table.h>
#include <zvec/core/framework/index_logger.h>
#include <zvec/ailego/logger/logger.h>
#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
} // namespace zvec

2
src/core/algorithm/hnsw/hnsw_streamer_entity.h
View File

@ -1099,4 +1099,4 @@ class HnswContiguousStreamerEntity : public HnswMmapStreamerEntity {
};
} // namespace core
} // namespace zvec
} // namespace zvec

1
src/db/CMakeLists.txt
View File

@ -46,6 +46,7 @@ cc_library(
FastPFOR
cppjieba
Arrow::arrow_static
Arrow::parquet_static
Arrow::arrow_compute
Arrow::arrow_dataset
Arrow::arrow_acero

2
src/db/common/global_resource.cc
View File

@ -30,4 +30,4 @@ void GlobalResource::initialize() {
});
}
} // namespace zvec
} // namespace zvec

1
src/db/index/CMakeLists.txt
View File

@ -25,6 +25,7 @@ cc_library(
rocksdb
core_interface
Arrow::arrow_static
Arrow::parquet_static
Arrow::arrow_compute
Arrow::arrow_dataset
cppjieba

12
src/db/index/storage/bufferpool_forward_store.cc
View File

@ -21,10 +21,10 @@
#include <arrow/result.h>
#include <arrow/status.h>
#include <parquet/arrow/reader.h>
#include <zvec/ailego/buffer/parquet_hash_table.h>
#include <zvec/ailego/logger/logger.h>
#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<std::string> &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<arrow::ChunkedArray> col_chunked_array =
buffer_handle.data();
if (!col_chunked_array) {
@ -317,9 +317,9 @@ ExecBatchPtr BufferPoolForwardStore::fetch(
std::vector<arrow::Datum> 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<arrow::ChunkedArray> col_chunked_array =
buffer_handle.data();
@ -379,4 +379,4 @@ RecordBatchReaderPtr BufferPoolForwardStore::scan(
physic_schema_, file_path_);
}
} // namespace zvec
} // namespace zvec

2
src/db/index/storage/bufferpool_forward_store.h
View File

@ -124,4 +124,4 @@ class BufferPoolForwardStore
std::vector<int64_t> row_group_row_nums_;
};
} // namespace zvec
} // namespace zvec

6
src/db/index/storage/lazy_record_batch_reader.h
View File

@ -16,8 +16,8 @@
#include <arrow/ipc/reader.h>
#include <parquet/arrow/reader.h>
#include <zvec/ailego/buffer/parquet_hash_table.h>
#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<arrow::ChunkedArray> col_chunked_array =
buffer_handle.data();
if (col_chunked_array) {

142
src/db/index/storage/parquet_buffer_pool.cc Normal file
View File

@ -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 <arrow/array/array_binary.h>
#include <arrow/io/file.h>
#include <arrow/ipc/reader.h>
#include <arrow/pretty_print.h>
#include <arrow/result.h>
#include <arrow/status.h>
#include <arrow/table.h>
#include <parquet/arrow/reader.h>
#include <zvec/core/framework/index_logger.h>
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<std::uint64_t>(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<arrow::io::RandomAccessFile> 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<parquet::arrow::FileReader> 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<arrow::Buffer> 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<arrow::ChunkedArray> ParquetBufferPool::retain(
ParquetBufferID buffer_id) {
return cache_.retain(buffer_id);
}
} // namespace zvec

127
src/include/zvec/ailego/buffer/parquet_hash_table.h → src/db/index/storage/parquet_buffer_pool.h
View File

@ -12,40 +12,37 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <sys/stat.h>
#include <atomic>
#include <chrono>
#include <cstdint>
#include <filesystem>
#include <memory>
#include <mutex>
#include <shared_mutex>
#include <string>
#include <vector>
#include <arrow/api.h>
#include <zvec/ailego/io/file.h>
#include <zvec/ailego/pattern/singleton.h>
#include "block_eviction_queue.h"
#include <zvec/ailego/buffer/external_cache.h>
namespace arrow {
class ChunkedArray;
class Array;
class DataType;
class Scalar;
template <typename T>
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<int>{}(1);
hash = hash ^ (std::hash<uint64_t>{}(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::ChunkedArray> arrow{nullptr};
// Guard original arrow buffers to prevent premature deletion
std::vector<std::shared_ptr<arrow::Buffer>> arrow_refs{};
size_t size;
alignas(64) std::atomic<int> ref_count{std::numeric_limits<int>::min()};
alignas(64) std::atomic<version_t> load_count{0};
};
class ParquetBufferLoader {
public:
using Value = std::shared_ptr<arrow::ChunkedArray>;
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::ChunkedArray> 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<ParquetBufferPool> 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<ParquetBufferID, ParquetBufferContext,
IDHash, IDEqual>;
arrow::Status acquire(ParquetBufferID buffer_id,
ParquetBufferContext &context);
ParquetBufferContextHandle acquire_buffer(ParquetBufferID buffer_id);
std::shared_ptr<arrow::ChunkedArray> set_block_acquired(
ParquetBufferID buffer_id);
std::shared_ptr<arrow::ChunkedArray> acquire(ParquetBufferID buffer_id);
std::shared_ptr<arrow::ChunkedArray> 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<ParquetBufferID, detail::ParquetBufferPayload,
detail::ParquetBufferLoader, ParquetBufferIDHash,
ParquetBufferIDEqual>;
ParquetBufferPool() = default;
std::shared_ptr<arrow::ChunkedArray> retain(ParquetBufferID buffer_id);
void release(ParquetBufferID buffer_id);
private:
Table table_;
std::shared_mutex table_mutex_;
Cache cache_;
};
} // namespace ailego
} // namespace zvec
} // namespace zvec

61
src/include/zvec/ailego/buffer/block_eviction_queue.h
View File

@ -22,7 +22,6 @@
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <limits>
#include <map>
#include <memory>
@ -34,7 +33,6 @@
#include <unordered_map>
#include <unordered_set>
#include <zvec/ailego/internal/platform.h>
#include <zvec/core/framework/index_logger.h>
#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<block_id_t, version_t> vector_block;
std::pair<ParquetBufferID, version_t> parquet_buffer_block;
VectorPageTable *page_table{nullptr};
eviction_key_t owner_key{0};
version_t version{0};
EvictableBlockOwner *owner{nullptr};
};
typedef moodycamel::ConcurrentQueue<BlockType> ConcurrentQueue;
@ -88,24 +83,24 @@ class BlockEvictionQueue {
// void clear_dead_node();
bool is_valid(VectorPageTable *page_table) {
std::shared_lock<std::shared_mutex> lock(valid_page_tables_mutex_);
return valid_page_tables_.find(page_table) != valid_page_tables_.end();
bool is_valid(EvictableBlockOwner *owner) {
std::shared_lock<std::shared_mutex> lock(valid_owners_mutex_);
return valid_owners_.find(owner) != valid_owners_.end();
}
void set_valid(VectorPageTable *page_table) {
std::unique_lock<std::shared_mutex> lock(valid_page_tables_mutex_);
valid_page_tables_.insert(page_table);
void set_valid(EvictableBlockOwner *owner) {
std::unique_lock<std::shared_mutex> lock(valid_owners_mutex_);
valid_owners_.insert(owner);
}
void set_invalid(VectorPageTable *page_table) {
std::unique_lock<std::shared_mutex> lock(valid_page_tables_mutex_);
valid_page_tables_.erase(page_table);
void set_invalid(EvictableBlockOwner *owner) {
std::unique_lock<std::shared_mutex> 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<ConcurrentQueue> evict_queues_;
std::unordered_set<VectorPageTable *> valid_page_tables_;
std::shared_mutex valid_page_tables_mutex_;
std::unordered_set<EvictableBlockOwner *> 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
} // namespace zvec

213
src/include/zvec/ailego/buffer/external_cache.h Normal file
View File

@ -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 <atomic>
#include <cstddef>
#include <limits>
#include <mutex>
#include <shared_mutex>
#include <unordered_map>
#include <utility>
#include <zvec/ailego/buffer/block_eviction_queue.h>
namespace zvec {
namespace ailego {
template <typename Key, typename Payload, typename Loader, typename Hash,
typename Equal>
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<std::shared_mutex> 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<std::shared_mutex> 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<std::shared_mutex> 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<std::shared_mutex> 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<std::shared_mutex> 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<std::shared_mutex> 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<int>::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<int> ref_count{std::numeric_limits<int>::min()};
alignas(64) std::atomic<version_t> 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<Key, Entry, Hash, Equal>;
Loader loader_{};
Table table_;
std::unordered_map<eviction_key_t, Key> owner_keys_;
eviction_key_t next_owner_key_{1};
std::shared_mutex mutex_;
};
} // namespace ailego
} // namespace zvec

11
src/include/zvec/ailego/buffer/vector_page_table.h
View File

@ -45,7 +45,7 @@ namespace ailego {
extern const size_t kVectorPageSize;
class VectorPageTable {
class VectorPageTable : public EvictableBlockOwner {
struct Entry {
std::atomic<int> ref_count;
std::atomic<bool> 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
} // namespace zvec

2
tests/ailego/CMakeLists.txt
View File

@ -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})

4
thirdparty/CMakeLists.txt vendored
View File

@ -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)