mirror of
https://github.com/milvus-io/milvus.git
synced 2026-10-11 22:49:53 +00:00
enhance: parallelize json stats build (#52120)
issue: #53765 JSON stats build benchmark — 500K rows per segment Concurrent segments || Non-parallel || 8 workers || Speedup 1 || 6.810 s || 1.641 s || 4.15× faster 4 || 7.140 s || 4.806 s || 1.49× faster 8 || 10.878 s || 9.658 s || 1.13× faster Setup: Release build, 12 CPUs, local filesystem, one warm-up followed by five measured runs; median values are reported. Data generation is excluded from the timing. Note: “Non-parallel” means each segment is built serially internally. Multiple segment builds are still launched concurrently in 8-segment cases. Signed-off-by: luzhang <[email protected]> Co-authored-by: luzhang <[email protected]>
This commit is contained in:
+5
-1
@@ -999,7 +999,7 @@ dataCoord:
|
||||
analyzeTaskSlotUsage: 65535 # slot usage of analyze task
|
||||
jsonShreddingMaxColumns: 1024 # the max number of columns to shred
|
||||
jsonShreddingRatioThreshold: 0.3 # the ratio threshold to shred
|
||||
jsonShreddingWriteBatchSize: 81920 # the batch size to write
|
||||
jsonShreddingWriteBatchSize: 81920 # the batch size to write; values must be positive, otherwise the default is used
|
||||
ip: # TCP/IP address of dataCoord. If not specified, use the first unicastable address
|
||||
port: 13333 # TCP port of dataCoord
|
||||
grpc:
|
||||
@@ -1079,6 +1079,10 @@ dataNode:
|
||||
storage:
|
||||
format: parquet # storage format for insert data, options: [parquet, vortex]
|
||||
deltalog: json # deltalog format, options: [json, parquet]
|
||||
jsonStats:
|
||||
parallelBuild:
|
||||
maxWorkers: 8 # Maximum number of process-wide JSON Stats Build workers. The effective value is capped by the cgroup-aware CPU count. Values less than or equal to 0 use the safe automatic limit min(CPU, 8).
|
||||
maxInflightBytes: 536870912 # Process-wide transient memory budget in bytes for materialized JSON Stats Build chunks. The budget is shared by concurrent builds. Lower values reduce peak memory at the cost of parallelism. A chunk larger than the budget is allowed to run exclusively to guarantee progress. Set to 0 to disable the limit.
|
||||
ip: # TCP/IP address of dataNode. If not specified, use the first unicastable address
|
||||
port: 21124 # TCP port of dataNode
|
||||
grpc:
|
||||
|
||||
@@ -10,7 +10,9 @@
|
||||
// or implied. See the License for the specific language governing permissions and limitations under the License
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstddef>
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
|
||||
#include "Executor.h"
|
||||
#include "folly/executors/CPUThreadPoolExecutor.h"
|
||||
@@ -20,6 +22,8 @@
|
||||
namespace milvus::futures {
|
||||
|
||||
const int kNumPriority = 3;
|
||||
const int kJsonStatsBuildNumPriority = 1;
|
||||
const size_t kJsonStatsBuildMinThreads = 1;
|
||||
|
||||
folly::CPUThreadPoolExecutor*
|
||||
getSearchCPUExecutor() {
|
||||
@@ -41,6 +45,18 @@ getLoadCPUExecutor() {
|
||||
return &executor;
|
||||
}
|
||||
|
||||
folly::CPUThreadPoolExecutor*
|
||||
getJsonStatsBuildExecutor() {
|
||||
auto max_thread_num = static_cast<size_t>(std::max(1, milvus::CPU_NUM));
|
||||
static folly::CPUThreadPoolExecutor executor(
|
||||
std::pair<size_t, size_t>{max_thread_num, kJsonStatsBuildMinThreads},
|
||||
folly::CPUThreadPoolExecutor::makeDefaultPriorityQueue(
|
||||
kJsonStatsBuildNumPriority),
|
||||
std::make_shared<folly::NamedThreadFactory>(
|
||||
"MILVUS_JSON_STATS_BUILD_"));
|
||||
return &executor;
|
||||
}
|
||||
|
||||
folly::CPUThreadPoolExecutor*
|
||||
getGlobalCPUExecutor() {
|
||||
return getSearchCPUExecutor();
|
||||
|
||||
@@ -33,4 +33,7 @@ getSearchCPUExecutor();
|
||||
folly::CPUThreadPoolExecutor*
|
||||
getLoadCPUExecutor();
|
||||
|
||||
folly::CPUThreadPoolExecutor*
|
||||
getJsonStatsBuildExecutor();
|
||||
|
||||
}; // namespace milvus::futures
|
||||
|
||||
@@ -16,6 +16,7 @@
|
||||
#include <stdlib.h>
|
||||
#include <chrono>
|
||||
#include <exception>
|
||||
#include <future>
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
#include <stdexcept>
|
||||
@@ -27,14 +28,65 @@
|
||||
|
||||
#include "common/EasyAssert.h"
|
||||
#include "common/common_type_c.h"
|
||||
#include "futures/Executor.h"
|
||||
#include "futures/Future.h"
|
||||
#include "futures/LeakyResult.h"
|
||||
#include "futures/Ready.h"
|
||||
#include "futures/future_c.h"
|
||||
#include "futures/future_c_types.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "index/json_stats/JsonStatsBuildMemoryBudget.h"
|
||||
|
||||
using namespace milvus::futures;
|
||||
|
||||
TEST(Futures, JsonStatsBuildExecutor) {
|
||||
auto* executor = getJsonStatsBuildExecutor();
|
||||
|
||||
ASSERT_NE(executor, nullptr);
|
||||
EXPECT_NE(executor, getSearchCPUExecutor());
|
||||
EXPECT_NE(executor, getLoadCPUExecutor());
|
||||
EXPECT_EQ(executor->getNumPriorities(), 1);
|
||||
EXPECT_GE(executor->numThreads(), 1);
|
||||
|
||||
std::promise<std::thread::id> promise;
|
||||
auto future = promise.get_future();
|
||||
executor->add(
|
||||
[&promise]() { promise.set_value(std::this_thread::get_id()); });
|
||||
|
||||
EXPECT_NE(future.get(), std::this_thread::get_id());
|
||||
}
|
||||
|
||||
TEST(Futures, ResizeJsonStatsBuildExecutor) {
|
||||
auto* executor = getJsonStatsBuildExecutor();
|
||||
auto original_thread_num = executor->numThreads();
|
||||
|
||||
auto status = executor_set_json_stats_build_thread_num(3);
|
||||
ASSERT_EQ(status.error_code, milvus::Success);
|
||||
EXPECT_EQ(executor->numThreads(), 3);
|
||||
|
||||
status = executor_set_json_stats_build_thread_num(0);
|
||||
ASSERT_EQ(status.error_code, milvus::Success);
|
||||
EXPECT_EQ(executor->numThreads(), 1);
|
||||
|
||||
status = executor_set_json_stats_build_thread_num(
|
||||
static_cast<int>(original_thread_num));
|
||||
ASSERT_EQ(status.error_code, milvus::Success);
|
||||
}
|
||||
|
||||
TEST(Futures, SetJsonStatsBuildMemoryBudget) {
|
||||
auto& budget = milvus::index::JsonStatsBuildMemoryBudget::GetInstance();
|
||||
auto original_capacity = budget.CapacityBytes();
|
||||
|
||||
auto status = executor_set_json_stats_build_max_inflight_bytes(1234);
|
||||
EXPECT_EQ(status.error_code, milvus::Success);
|
||||
EXPECT_EQ(budget.CapacityBytes(), 1234);
|
||||
|
||||
status = executor_set_json_stats_build_max_inflight_bytes(
|
||||
static_cast<int64_t>(original_capacity));
|
||||
EXPECT_EQ(status.error_code, milvus::Success);
|
||||
EXPECT_EQ(budget.CapacityBytes(), original_capacity);
|
||||
}
|
||||
|
||||
TEST(Futures, LeakyResult) {
|
||||
{
|
||||
LeakyResult<int> leaky_result;
|
||||
|
||||
@@ -9,6 +9,10 @@
|
||||
// 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 <algorithm>
|
||||
#include <exception>
|
||||
|
||||
#include "common/EasyAssert.h"
|
||||
#include "Executor.h"
|
||||
#include "Future.h"
|
||||
#include "folly/executors/CPUThreadPoolExecutor.h"
|
||||
@@ -16,6 +20,7 @@
|
||||
#include "common/CGoCatch.h"
|
||||
#include "futures/future_c_types.h"
|
||||
#include "glog/logging.h"
|
||||
#include "index/json_stats/JsonStatsBuildMemoryBudget.h"
|
||||
#include "log/Log.h"
|
||||
#include "monitor/Monitor.h"
|
||||
#include "prometheus/gauge.h"
|
||||
@@ -82,3 +87,34 @@ executor_set_load_thread_num(int thread_num) {
|
||||
}
|
||||
CGO_CATCH_AND_LOG("executor_set_load_thread_num")
|
||||
}
|
||||
|
||||
extern "C" CStatus
|
||||
executor_set_json_stats_build_thread_num(int thread_num) {
|
||||
auto normalized_thread_num = std::max(1, thread_num);
|
||||
try {
|
||||
milvus::futures::getJsonStatsBuildExecutor()->setNumThreads(
|
||||
normalized_thread_num);
|
||||
LOG_INFO(
|
||||
"future executor setup json stats build cpu executor with thread "
|
||||
"num: {}",
|
||||
normalized_thread_num);
|
||||
return milvus::SuccessCStatus();
|
||||
}
|
||||
CGO_CATCH_AND_RETURN_CSTATUS
|
||||
}
|
||||
|
||||
extern "C" CStatus
|
||||
executor_set_json_stats_build_max_inflight_bytes(int64_t max_inflight_bytes) {
|
||||
try {
|
||||
auto normalized_max_inflight_bytes =
|
||||
static_cast<size_t>(std::max<int64_t>(0, max_inflight_bytes));
|
||||
milvus::index::JsonStatsBuildMemoryBudget::GetInstance()
|
||||
.SetCapacityBytes(normalized_max_inflight_bytes);
|
||||
LOG_INFO(
|
||||
"json stats build transient memory budget setup with max inflight "
|
||||
"bytes: {}",
|
||||
normalized_max_inflight_bytes);
|
||||
return milvus::SuccessCStatus();
|
||||
}
|
||||
CGO_CATCH_AND_RETURN_CSTATUS
|
||||
}
|
||||
|
||||
@@ -14,6 +14,7 @@
|
||||
#include "common/common_type_c.h"
|
||||
#include "future_c_types.h"
|
||||
#include <stdbool.h>
|
||||
#include <stdint.h>
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
@@ -46,6 +47,12 @@ executor_set_search_thread_num(int thread_num);
|
||||
void
|
||||
executor_set_load_thread_num(int thread_num);
|
||||
|
||||
CStatus
|
||||
executor_set_json_stats_build_thread_num(int thread_num);
|
||||
|
||||
CStatus
|
||||
executor_set_json_stats_build_max_inflight_bytes(int64_t max_inflight_bytes);
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
|
||||
@@ -14,17 +14,21 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include <nlohmann/json.hpp>
|
||||
#include <string.h>
|
||||
#include <algorithm>
|
||||
#include <atomic>
|
||||
#include <chrono>
|
||||
#include <cstdint>
|
||||
#include <exception>
|
||||
#include <filesystem>
|
||||
#include <initializer_list>
|
||||
#include <iosfwd>
|
||||
#include <map>
|
||||
#include <memory>
|
||||
#include <string.h>
|
||||
#include <unordered_set>
|
||||
#include <utility>
|
||||
#include <variant>
|
||||
#include "segcore/default_fs.h"
|
||||
#include <vector>
|
||||
|
||||
#include "NamedType/named_type_impl.hpp"
|
||||
#include "NamedType/underlying_functionalities.hpp"
|
||||
@@ -38,12 +42,16 @@
|
||||
#include "common/GroupChunk.h"
|
||||
#include "common/Json.h"
|
||||
#include "common/Tracer.h"
|
||||
#include "common/Utils.h"
|
||||
#include "common/jsmn.h"
|
||||
#include "fmt/core.h"
|
||||
#include "folly/ScopeGuard.h"
|
||||
#include "folly/coro/BlockingWait.h"
|
||||
#include "folly/futures/Future.h"
|
||||
#include "futures/Executor.h"
|
||||
#include "index/Utils.h"
|
||||
#include "index/json_stats/JsonKeyStats.h"
|
||||
#include "index/json_stats/JsonStatsBuildMemoryBudget.h"
|
||||
#include "index/json_stats/bson_builder.h"
|
||||
#include "index/json_stats/parquet_writer.h"
|
||||
#include "milvus-storage/common/config.h"
|
||||
@@ -57,15 +65,17 @@
|
||||
#include "milvus-storage/reader.h"
|
||||
#include "mmap/ChunkedColumnGroup.h"
|
||||
#include "mmap/Types.h"
|
||||
#include "nlohmann/json.hpp"
|
||||
#include "nlohmann/detail/iterators/iteration_proxy.hpp"
|
||||
#include "nlohmann/json_fwd.hpp"
|
||||
#include "parquet/metadata.h"
|
||||
#include "segcore/storagev1translator/BsonInvertedIndexTranslator.h"
|
||||
#include "segcore/ChunkedSegmentSealedImpl.h"
|
||||
#include "segcore/storagev2translator/AsyncChunkReader.h"
|
||||
#include "segcore/storagev2translator/ManifestGroupTranslator.h"
|
||||
#include "segcore/storagev2translator/StorageV2Config.h"
|
||||
#include "segcore/Utils.h"
|
||||
#include "segcore/default_fs.h"
|
||||
#include "segcore/storagev1translator/BsonInvertedIndexTranslator.h"
|
||||
#include "storage/DiskFileManagerImpl.h"
|
||||
#include "storage/FileManager.h"
|
||||
#include "storage/LocalChunkManager.h"
|
||||
@@ -84,6 +94,88 @@ namespace {
|
||||
|
||||
// Reader::create() exposes this synthetic JSON-stats group at index zero.
|
||||
constexpr int64_t kJsonStatsReaderColumnGroupIndex = 0;
|
||||
constexpr int64_t kJsonStatsRowsPerRange = 16 * 1024;
|
||||
constexpr size_t kJsonStatsMaterializeMinReservationBytes = 64 * 1024;
|
||||
constexpr size_t kJsonStatsMaterializeInputExpansionFactor = 2;
|
||||
constexpr size_t kJsonStatsMaterializeRowOverheadBytes = 32;
|
||||
constexpr size_t kJsonStatsMaterializeColumnOverheadBytes = 1024;
|
||||
|
||||
struct CollectKeyInfoResult {
|
||||
std::map<JsonKey, KeyStatsInfo> infos;
|
||||
};
|
||||
|
||||
template <typename Func>
|
||||
void
|
||||
ForEachJsonStatsRow(const JsonStatsRowRange& range, Func&& func) {
|
||||
int64_t visited_rows = 0;
|
||||
for (const auto& slice : range.slices) {
|
||||
AssertInfo(slice.data != nullptr,
|
||||
"json stats row range contains null field data");
|
||||
AssertInfo(slice.row_count > 0,
|
||||
"json stats row range contains an empty slice");
|
||||
AssertInfo(slice.local_begin >= 0,
|
||||
"json stats field data slice has negative local begin: {}",
|
||||
slice.local_begin);
|
||||
AssertInfo(
|
||||
slice.local_begin <= slice.data->get_num_rows() - slice.row_count,
|
||||
"json stats field data slice [{}, {}) exceeds field data rows {}",
|
||||
slice.local_begin,
|
||||
slice.local_begin + slice.row_count,
|
||||
slice.data->get_num_rows());
|
||||
const auto slice_global_begin = range.global_begin + visited_rows;
|
||||
|
||||
for (int64_t offset = 0; offset < slice.row_count; ++offset) {
|
||||
func(slice.data,
|
||||
slice.local_begin + offset,
|
||||
slice_global_begin + offset);
|
||||
}
|
||||
visited_rows += slice.row_count;
|
||||
}
|
||||
AssertInfo(visited_rows == range.row_count,
|
||||
"json stats row range contains {} rows, expected {}",
|
||||
visited_rows,
|
||||
range.row_count);
|
||||
}
|
||||
|
||||
size_t
|
||||
BytesForBits(size_t bit_count) {
|
||||
return SaturatingAdd(bit_count / 8,
|
||||
bit_count % 8 == 0 ? size_t{0} : size_t{1});
|
||||
}
|
||||
|
||||
size_t
|
||||
EstimateArrowFixedBufferBytes(const arrow::Schema& schema, size_t row_count) {
|
||||
const auto validity_bytes = BytesForBits(row_count);
|
||||
size_t buffer_bytes = 0;
|
||||
for (const auto& field : schema.fields()) {
|
||||
// Every materialized column can contain nulls. Account for its validity
|
||||
// bitmap even when a dense finished array can omit that buffer.
|
||||
buffer_bytes = SaturatingAdd(buffer_bytes, validity_bytes);
|
||||
|
||||
const auto& type = *field->type();
|
||||
if (type.id() == arrow::Type::STRING ||
|
||||
type.id() == arrow::Type::BINARY) {
|
||||
// The variable payload is covered by the input-size estimate. This
|
||||
// accounts for the offsets that are allocated even for null rows.
|
||||
const auto offset_count = SaturatingAdd(row_count, size_t{1});
|
||||
buffer_bytes = SaturatingAdd(
|
||||
buffer_bytes,
|
||||
SaturatingMultiply(offset_count,
|
||||
sizeof(arrow::BinaryType::offset_type)));
|
||||
continue;
|
||||
}
|
||||
|
||||
const auto bit_width = type.bit_width();
|
||||
AssertInfo(bit_width > 0,
|
||||
"unsupported json stats Arrow field type {}",
|
||||
type.ToString());
|
||||
buffer_bytes =
|
||||
SaturatingAdd(buffer_bytes,
|
||||
BytesForBits(SaturatingMultiply(
|
||||
row_count, static_cast<size_t>(bit_width))));
|
||||
}
|
||||
return buffer_bytes;
|
||||
}
|
||||
|
||||
struct JsonStatsParquetMetadata {
|
||||
std::shared_ptr<arrow::Schema> schema;
|
||||
@@ -205,6 +297,93 @@ GetJsonStatsFieldsFromSchema(const std::shared_ptr<arrow::Schema>& schema) {
|
||||
|
||||
} // namespace
|
||||
|
||||
std::vector<JsonStatsRowRange>
|
||||
CreateJsonStatsRowRanges(const std::vector<FieldDataPtr>& field_datas,
|
||||
int64_t max_rows_per_range) {
|
||||
AssertInfo(max_rows_per_range > 0,
|
||||
"json stats max rows per range must be positive, got {}",
|
||||
max_rows_per_range);
|
||||
|
||||
std::vector<JsonStatsRowRange> row_ranges;
|
||||
JsonStatsRowRange current_range{};
|
||||
int64_t global_begin = 0;
|
||||
|
||||
for (const auto& data : field_datas) {
|
||||
AssertInfo(data != nullptr,
|
||||
"cannot create json stats row range from null field data");
|
||||
const auto data_rows = data->get_num_rows();
|
||||
AssertInfo(data_rows >= 0,
|
||||
"json stats field data has negative row count: {}",
|
||||
data_rows);
|
||||
|
||||
int64_t local_begin = 0;
|
||||
while (local_begin < data_rows) {
|
||||
if (current_range.row_count == 0) {
|
||||
current_range.global_begin = global_begin + local_begin;
|
||||
}
|
||||
|
||||
const auto remaining_rows =
|
||||
max_rows_per_range - current_range.row_count;
|
||||
const auto slice_rows =
|
||||
std::min(data_rows - local_begin, remaining_rows);
|
||||
current_range.slices.push_back(JsonStatsFieldDataSlice{
|
||||
data,
|
||||
local_begin,
|
||||
slice_rows,
|
||||
});
|
||||
current_range.row_count += slice_rows;
|
||||
local_begin += slice_rows;
|
||||
|
||||
if (current_range.row_count == max_rows_per_range) {
|
||||
row_ranges.push_back(std::move(current_range));
|
||||
current_range = JsonStatsRowRange{};
|
||||
}
|
||||
}
|
||||
global_begin += data_rows;
|
||||
}
|
||||
|
||||
if (current_range.row_count > 0) {
|
||||
row_ranges.push_back(std::move(current_range));
|
||||
}
|
||||
return row_ranges;
|
||||
}
|
||||
|
||||
size_t
|
||||
JsonKeyStats::EstimateJsonStatsMaterializeReservationBytes(
|
||||
const JsonStatsRowRange& range, const arrow::Schema& schema) {
|
||||
size_t input_bytes = 0;
|
||||
ForEachJsonStatsRow(
|
||||
range,
|
||||
[&](const FieldDataPtr& data,
|
||||
int64_t local_row,
|
||||
int64_t /* global_row */) {
|
||||
if (data->IsNullable() && !data->is_valid(local_row)) {
|
||||
return;
|
||||
}
|
||||
auto raw_value = data->RawValue(local_row);
|
||||
if (raw_value == nullptr) {
|
||||
return;
|
||||
}
|
||||
const auto* json = static_cast<const milvus::Json*>(raw_value);
|
||||
input_bytes = SaturatingAdd(input_bytes, json->data().size());
|
||||
});
|
||||
|
||||
const auto row_count = static_cast<size_t>(range.row_count);
|
||||
auto reservation_bytes = SaturatingMultiply(
|
||||
input_bytes, kJsonStatsMaterializeInputExpansionFactor);
|
||||
reservation_bytes = SaturatingAdd(
|
||||
reservation_bytes,
|
||||
SaturatingMultiply(row_count, kJsonStatsMaterializeRowOverheadBytes));
|
||||
reservation_bytes = SaturatingAdd(
|
||||
reservation_bytes,
|
||||
SaturatingMultiply(static_cast<size_t>(schema.num_fields()),
|
||||
kJsonStatsMaterializeColumnOverheadBytes));
|
||||
reservation_bytes = SaturatingAdd(
|
||||
reservation_bytes, EstimateArrowFixedBufferBytes(schema, row_count));
|
||||
return std::max(kJsonStatsMaterializeMinReservationBytes,
|
||||
reservation_bytes);
|
||||
}
|
||||
|
||||
JsonKeyStats::JsonKeyStats(const storage::FileManagerContext& ctx,
|
||||
bool is_load,
|
||||
int64_t json_stats_max_shredding_columns,
|
||||
@@ -221,7 +400,11 @@ JsonKeyStats::JsonKeyStats(const storage::FileManagerContext& ctx,
|
||||
std::make_shared<milvus::storage::MemFileManagerImpl>(ctx);
|
||||
disk_file_manager_ =
|
||||
std::make_shared<milvus::storage::DiskFileManagerImpl>(ctx);
|
||||
write_batch_size_ = json_stats_write_batch_size;
|
||||
// Older coordinators can send values that bypass the Go config formatter.
|
||||
// Normalize before choosing row ranges or passing the value to the writer.
|
||||
write_batch_size_ = json_stats_write_batch_size > 0
|
||||
? json_stats_write_batch_size
|
||||
: kDefaultWriteBatchSize;
|
||||
max_shredding_columns_ = json_stats_max_shredding_columns;
|
||||
shredding_ratio_threshold_ = json_stats_shredding_ratio_threshold;
|
||||
LOG_INFO(
|
||||
@@ -445,6 +628,100 @@ JsonKeyStats::CollectKeyInfo(const std::vector<FieldDataPtr>& field_datas,
|
||||
return infos;
|
||||
}
|
||||
|
||||
std::map<JsonKey, KeyStatsInfo>
|
||||
JsonKeyStats::CollectKeyInfo(const std::vector<JsonStatsRowRange>& row_ranges,
|
||||
bool nullable) {
|
||||
std::map<JsonKey, KeyStatsInfo> infos;
|
||||
if (row_ranges.empty()) {
|
||||
num_rows_ = 0;
|
||||
return infos;
|
||||
}
|
||||
|
||||
auto* executor = milvus::futures::getJsonStatsBuildExecutor();
|
||||
const auto max_inflight = std::min(
|
||||
row_ranges.size(), std::max<size_t>(1, executor->numThreads()));
|
||||
|
||||
// Futures are kept in range-index order. Only max_inflight tasks are active
|
||||
// at a time, so completed local maps cannot accumulate for every range.
|
||||
std::vector<folly::Future<CollectKeyInfoResult>> futures;
|
||||
futures.reserve(row_ranges.size());
|
||||
size_t next_submit = 0;
|
||||
size_t next_consume = 0;
|
||||
std::exception_ptr first_exception;
|
||||
|
||||
auto submit_next = [&]() {
|
||||
try {
|
||||
auto range = row_ranges[next_submit];
|
||||
auto future = folly::via(
|
||||
executor, [this, range = std::move(range), nullable]() {
|
||||
std::map<JsonKey, KeyStatsInfo> local_infos;
|
||||
ForEachJsonStatsRow(
|
||||
range,
|
||||
[&](const FieldDataPtr& data,
|
||||
int64_t local_row,
|
||||
int64_t /* global_row */) {
|
||||
if ((nullable || data->IsNullable()) &&
|
||||
!data->is_valid(local_row)) {
|
||||
return;
|
||||
}
|
||||
auto json_str = static_cast<const milvus::Json*>(
|
||||
data->RawValue(local_row))
|
||||
->data();
|
||||
CollectSingleJsonStatsInfo(json_str, local_infos);
|
||||
});
|
||||
return CollectKeyInfoResult{std::move(local_infos)};
|
||||
});
|
||||
futures.emplace_back(std::move(future));
|
||||
++next_submit;
|
||||
} catch (...) {
|
||||
if (first_exception == nullptr) {
|
||||
first_exception = std::current_exception();
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
while (next_submit < max_inflight && first_exception == nullptr) {
|
||||
submit_next();
|
||||
}
|
||||
|
||||
int64_t num_rows = 0;
|
||||
while (next_consume < futures.size()) {
|
||||
try {
|
||||
auto result = std::move(futures[next_consume]).get();
|
||||
if (first_exception == nullptr) {
|
||||
const auto& range = row_ranges[next_consume];
|
||||
AssertInfo(
|
||||
range.global_begin == num_rows,
|
||||
"json stats row range starts at global row {}, expected {}",
|
||||
range.global_begin,
|
||||
num_rows);
|
||||
for (const auto& [key, local_info] : result.infos) {
|
||||
infos[key].hit_row_num_ += local_info.hit_row_num_;
|
||||
}
|
||||
num_rows += range.row_count;
|
||||
}
|
||||
} catch (...) {
|
||||
if (first_exception == nullptr) {
|
||||
first_exception = std::current_exception();
|
||||
}
|
||||
}
|
||||
++next_consume;
|
||||
|
||||
if (first_exception == nullptr && next_submit < row_ranges.size()) {
|
||||
submit_next();
|
||||
}
|
||||
}
|
||||
|
||||
// Every submitted future has been consumed before an error escapes. This
|
||||
// keeps `this` and the FieldData referenced by worker lambdas alive.
|
||||
if (first_exception != nullptr) {
|
||||
std::rethrow_exception(first_exception);
|
||||
}
|
||||
|
||||
num_rows_ = num_rows;
|
||||
return infos;
|
||||
}
|
||||
|
||||
std::map<JsonKey, JsonKeyLayoutType>
|
||||
JsonKeyStats::ClassifyJsonKeyLayoutType(
|
||||
const std::map<JsonKey, KeyStatsInfo>& infos) {
|
||||
@@ -551,7 +828,7 @@ void
|
||||
JsonKeyStats::AddKeyStats(const std::vector<std::string>& path,
|
||||
JSONType type,
|
||||
const std::string& value,
|
||||
std::map<JsonKey, std::string>& values) {
|
||||
std::map<JsonKey, std::string>& values) const {
|
||||
auto path_str = JsonPointer(path);
|
||||
auto key = JsonKey(path_str, type);
|
||||
values[key] = value;
|
||||
@@ -563,7 +840,7 @@ JsonKeyStats::TraverseJsonForBuildStats(
|
||||
jsmntok* tokens,
|
||||
int& index,
|
||||
std::vector<std::string>& path,
|
||||
std::map<JsonKey, std::string>& values) {
|
||||
std::map<JsonKey, std::string>& values) const {
|
||||
jsmntok current = tokens[0];
|
||||
if (!(current.type != JSMN_UNDEFINED)) {
|
||||
ThrowInfo(ErrorCode::DataFormatBroken,
|
||||
@@ -604,15 +881,7 @@ JsonKeyStats::TraverseJsonForBuildStats(
|
||||
index = j;
|
||||
} else if (current.type == JSMN_PRIMITIVE) {
|
||||
std::string value(json + current.start, current.end - current.start);
|
||||
JSONType type;
|
||||
try {
|
||||
type = getType(value);
|
||||
} catch (const std::exception& e) {
|
||||
ThrowInfo(ErrorCode::DataFormatBroken,
|
||||
"failed to get json type for value: {} with error: {}",
|
||||
value,
|
||||
e.what());
|
||||
}
|
||||
auto type = getType(value);
|
||||
|
||||
if (type == JSONType::INT64) {
|
||||
AddKeyStats(path, JSONType::INT64, value, values);
|
||||
@@ -656,26 +925,9 @@ JsonKeyStats::TraverseJsonForBuildStats(
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
JsonKeyStats::BuildKeyStatsForNullRow() {
|
||||
// add empty value for column keys that not hit
|
||||
for (const auto& key : column_keys_) {
|
||||
parquet_writer_->AppendValue(key.ToColumnName(), "");
|
||||
}
|
||||
|
||||
// add null bson to shared column
|
||||
BsonDocument null_doc;
|
||||
parquet_writer_->AppendSharedRow(null_doc.data(), null_doc.length());
|
||||
|
||||
parquet_writer_->AddCurrentRow();
|
||||
}
|
||||
|
||||
void
|
||||
JsonKeyStats::BuildKeyStatsForRow(std::string_view json_str, uint32_t row_id) {
|
||||
LOG_TRACE("build key stats for row {} with json {} for segment {}",
|
||||
row_id,
|
||||
json_str,
|
||||
segment_id_);
|
||||
bool
|
||||
JsonKeyStats::ParseJsonForBuildStats(
|
||||
std::string_view json_str, std::map<JsonKey, std::string>& values) const {
|
||||
jsmn_parser parser;
|
||||
jsmn_init(&parser);
|
||||
|
||||
@@ -691,30 +943,56 @@ JsonKeyStats::BuildKeyStatsForRow(std::string_view json_str, uint32_t row_id) {
|
||||
token_capacity);
|
||||
if (r < 0) {
|
||||
if (r == JSMN_ERROR_NOMEM) {
|
||||
// Reallocate tokens array if not enough space
|
||||
token_capacity *= 2;
|
||||
tokens.resize(token_capacity);
|
||||
continue;
|
||||
} else {
|
||||
ThrowInfo(ErrorCode::DataFormatBroken,
|
||||
"Failed to parse Json: {}, error: {}",
|
||||
json_str,
|
||||
int(r));
|
||||
}
|
||||
ThrowInfo(ErrorCode::DataFormatBroken,
|
||||
"Failed to parse Json: {}, error: {}",
|
||||
json_str,
|
||||
int(r));
|
||||
}
|
||||
num_tokens = r;
|
||||
break;
|
||||
}
|
||||
|
||||
if (num_tokens == 0) {
|
||||
return;
|
||||
return false;
|
||||
}
|
||||
|
||||
int index = 0;
|
||||
std::vector<std::string> paths;
|
||||
std::map<JsonKey, std::string> values;
|
||||
TraverseJsonForBuildStats(
|
||||
json_str.data(), tokens.data(), index, paths, values);
|
||||
return true;
|
||||
}
|
||||
|
||||
void
|
||||
JsonKeyStats::BuildKeyStatsForNullRow() {
|
||||
// add empty value for column keys that not hit
|
||||
for (const auto& key : column_keys_) {
|
||||
parquet_writer_->AppendValue(key.ToColumnName(), "");
|
||||
}
|
||||
|
||||
// add an empty BSON document to the shared column
|
||||
BsonDocument null_doc;
|
||||
parquet_writer_->AppendSharedRow(null_doc.data(), null_doc.length());
|
||||
|
||||
parquet_writer_->AddCurrentRow();
|
||||
}
|
||||
|
||||
void
|
||||
JsonKeyStats::BuildKeyStatsForRow(std::string_view json_str, uint32_t row_id) {
|
||||
LOG_TRACE("build key stats for row {} with json {} for segment {}",
|
||||
row_id,
|
||||
json_str,
|
||||
segment_id_);
|
||||
std::map<JsonKey, std::string> values;
|
||||
if (!ParseJsonForBuildStats(json_str, values)) {
|
||||
BuildKeyStatsForNullRow();
|
||||
return;
|
||||
}
|
||||
|
||||
DomNode root;
|
||||
std::set<JsonKey> hit_keys;
|
||||
for (const auto& [key, value] : values) {
|
||||
@@ -769,29 +1047,430 @@ JsonKeyStats::BuildKeyStatsForRow(std::string_view json_str, uint32_t row_id) {
|
||||
}
|
||||
|
||||
void
|
||||
JsonKeyStats::BuildKeyStats(const std::vector<FieldDataPtr>& field_datas,
|
||||
JsonKeyStats::BuildKeyStats(const std::vector<JsonStatsRowRange>& row_ranges,
|
||||
bool nullable) {
|
||||
uint32_t row_id = 0;
|
||||
for (const auto& data : field_datas) {
|
||||
auto n = data->get_num_rows();
|
||||
for (uint32_t i = 0; i < n; i++) {
|
||||
if ((nullable || data->IsNullable()) && !data->is_valid(i)) {
|
||||
BuildKeyStatsForNullRow();
|
||||
} else {
|
||||
auto json_str =
|
||||
static_cast<const milvus::Json*>(data->RawValue(i))->data();
|
||||
|
||||
// some situations, such as empty json string,
|
||||
// should be handled as null row
|
||||
if (json_str.empty()) {
|
||||
int64_t processed_rows = 0;
|
||||
for (const auto& range : row_ranges) {
|
||||
AssertInfo(range.global_begin == processed_rows,
|
||||
"json stats row range starts at global row {}, expected {}",
|
||||
range.global_begin,
|
||||
processed_rows);
|
||||
ForEachJsonStatsRow(
|
||||
range,
|
||||
[&](const FieldDataPtr& data,
|
||||
int64_t local_row,
|
||||
int64_t global_row) {
|
||||
if ((nullable || data->IsNullable()) &&
|
||||
!data->is_valid(local_row)) {
|
||||
BuildKeyStatsForNullRow();
|
||||
} else {
|
||||
BuildKeyStatsForRow(json_str, row_id);
|
||||
auto json_str = static_cast<const milvus::Json*>(
|
||||
data->RawValue(local_row))
|
||||
->data();
|
||||
|
||||
// some situations, such as empty json string,
|
||||
// should be handled as null row
|
||||
if (json_str.empty()) {
|
||||
BuildKeyStatsForNullRow();
|
||||
} else {
|
||||
BuildKeyStatsForRow(json_str,
|
||||
static_cast<uint32_t>(global_row));
|
||||
}
|
||||
}
|
||||
}
|
||||
row_id++;
|
||||
});
|
||||
processed_rows += range.row_count;
|
||||
}
|
||||
AssertInfo(processed_rows == num_rows_,
|
||||
"materialized {} json stats rows, expected {}",
|
||||
processed_rows,
|
||||
num_rows_);
|
||||
}
|
||||
|
||||
JsonKeyStats::MaterializedChunk
|
||||
JsonKeyStats::MaterializeKeyStatsRange(
|
||||
const JsonStatsRowRange& range,
|
||||
bool nullable,
|
||||
const std::shared_ptr<arrow::Schema>& schema) const {
|
||||
AssertInfo(schema != nullptr,
|
||||
"json stats materialize schema must not be null");
|
||||
|
||||
auto [builders, builders_map] = CreateArrowBuildersForColumns(column_keys_);
|
||||
AssertInfo(
|
||||
builders.size() == static_cast<size_t>(schema->num_fields()),
|
||||
"json stats materialize builder count {} does not match schema field "
|
||||
"count {}",
|
||||
builders.size(),
|
||||
schema->num_fields());
|
||||
// Match the pre-dispatch fixed-buffer estimate and avoid geometric
|
||||
// over-allocation while appending a partially filled final range.
|
||||
for (const auto& builder : builders) {
|
||||
auto status = builder->Reserve(range.row_count);
|
||||
if (!status.ok()) {
|
||||
ThrowInfo(milvus::storage::ArrowStatusToErrorCode(status),
|
||||
"failed to reserve json stats builder for {} rows: {}",
|
||||
range.row_count,
|
||||
status.ToString());
|
||||
}
|
||||
}
|
||||
|
||||
std::map<std::string, std::vector<int64_t>> bson_postings;
|
||||
auto append_null_row = [&]() {
|
||||
for (size_t i = 0; i + 1 < builders.size(); ++i) {
|
||||
auto status = builders[i]->AppendNull();
|
||||
if (!status.ok()) {
|
||||
ThrowInfo(milvus::storage::ArrowStatusToErrorCode(status),
|
||||
"failed to append null json stats value: {}",
|
||||
status.ToString());
|
||||
}
|
||||
}
|
||||
|
||||
BsonDocument null_doc;
|
||||
auto shared_builder =
|
||||
std::static_pointer_cast<arrow::BinaryBuilder>(builders.back());
|
||||
auto status =
|
||||
shared_builder->Append(null_doc.data(), null_doc.length());
|
||||
if (!status.ok()) {
|
||||
ThrowInfo(milvus::storage::ArrowStatusToErrorCode(status),
|
||||
"failed to append empty shared json stats value: {}",
|
||||
status.ToString());
|
||||
}
|
||||
};
|
||||
|
||||
ForEachJsonStatsRow(
|
||||
range,
|
||||
[&](const FieldDataPtr& data, int64_t local_row, int64_t global_row) {
|
||||
if ((nullable || data->IsNullable()) &&
|
||||
!data->is_valid(local_row)) {
|
||||
append_null_row();
|
||||
return;
|
||||
}
|
||||
|
||||
auto json_str =
|
||||
static_cast<const milvus::Json*>(data->RawValue(local_row))
|
||||
->data();
|
||||
if (json_str.empty()) {
|
||||
append_null_row();
|
||||
return;
|
||||
}
|
||||
|
||||
LOG_TRACE(
|
||||
"materialize key stats for row {} with json {} for "
|
||||
"segment {}",
|
||||
global_row,
|
||||
json_str,
|
||||
segment_id_);
|
||||
std::map<JsonKey, std::string> values;
|
||||
if (!ParseJsonForBuildStats(json_str, values)) {
|
||||
append_null_row();
|
||||
return;
|
||||
}
|
||||
|
||||
DomNode root;
|
||||
std::set<JsonKey> hit_keys;
|
||||
for (const auto& [key, value] : values) {
|
||||
auto key_type = key_types_.find(key);
|
||||
AssertInfo(key_type != key_types_.end(),
|
||||
"key {} not found in key types",
|
||||
key.key_);
|
||||
if (key_type->second == JsonKeyLayoutType::SHARED) {
|
||||
auto path_vec = ParseJsonPointerPath(key.key_);
|
||||
BsonBuilder::AppendToDom(root, path_vec, value, key.type_);
|
||||
} else {
|
||||
auto builder = builders_map.find(key.ToColumnName());
|
||||
AssertInfo(builder != builders_map.end(),
|
||||
"builder for key {} not found",
|
||||
key.ToColumnName());
|
||||
|
||||
arrow::Status status;
|
||||
if (key.type_ == JSONType::ARRAY) {
|
||||
auto bson_bytes =
|
||||
BuildBsonArrayBytesFromJsonString(value);
|
||||
status = AppendJsonStatsValueToBuilder(
|
||||
std::string(reinterpret_cast<const char*>(
|
||||
bson_bytes.data()),
|
||||
bson_bytes.size()),
|
||||
builder->second);
|
||||
} else {
|
||||
status = AppendJsonStatsValueToBuilder(value,
|
||||
builder->second);
|
||||
}
|
||||
if (!status.ok()) {
|
||||
ThrowInfo(
|
||||
milvus::storage::ArrowStatusToErrorCode(status),
|
||||
"failed to append json stats value for key {}: {}",
|
||||
key.ToColumnName(),
|
||||
status.ToString());
|
||||
}
|
||||
}
|
||||
hit_keys.insert(key);
|
||||
}
|
||||
|
||||
for (const auto& key : column_keys_) {
|
||||
if (hit_keys.find(key) != hit_keys.end()) {
|
||||
continue;
|
||||
}
|
||||
auto builder = builders_map.find(key.ToColumnName());
|
||||
AssertInfo(builder != builders_map.end(),
|
||||
"builder for key {} not found",
|
||||
key.ToColumnName());
|
||||
auto status = builder->second->AppendNull();
|
||||
if (!status.ok()) {
|
||||
ThrowInfo(milvus::storage::ArrowStatusToErrorCode(status),
|
||||
"failed to append null json stats value for key "
|
||||
"{}: {}",
|
||||
key.ToColumnName(),
|
||||
status.ToString());
|
||||
}
|
||||
}
|
||||
|
||||
BsonDocument final_doc;
|
||||
BsonBuilder::ConvertDomToBson(root, final_doc.get());
|
||||
auto key_offsets = BsonBuilder::ExtractBsonKeyOffsets(
|
||||
final_doc.data(), final_doc.length());
|
||||
for (const auto& [key, offset] : key_offsets) {
|
||||
bson_postings[key].push_back(EncodeInvertedIndexValue(
|
||||
static_cast<uint32_t>(global_row), offset));
|
||||
}
|
||||
|
||||
auto shared_builder =
|
||||
std::static_pointer_cast<arrow::BinaryBuilder>(builders.back());
|
||||
auto status =
|
||||
shared_builder->Append(final_doc.data(), final_doc.length());
|
||||
if (!status.ok()) {
|
||||
ThrowInfo(milvus::storage::ArrowStatusToErrorCode(status),
|
||||
"failed to append shared json stats value: {}",
|
||||
status.ToString());
|
||||
}
|
||||
});
|
||||
|
||||
std::vector<std::shared_ptr<arrow::Array>> arrays;
|
||||
arrays.reserve(builders.size());
|
||||
for (const auto& builder : builders) {
|
||||
AssertInfo(builder->length() == range.row_count,
|
||||
"json stats materialize builder contains {} rows, "
|
||||
"expected {}",
|
||||
builder->length(),
|
||||
range.row_count);
|
||||
std::shared_ptr<arrow::Array> array;
|
||||
auto status = builder->Finish(&array);
|
||||
if (!status.ok()) {
|
||||
ThrowInfo(milvus::storage::ArrowStatusToErrorCode(status),
|
||||
"failed to finish json stats materialize builder: {}",
|
||||
status.ToString());
|
||||
}
|
||||
arrays.push_back(std::move(array));
|
||||
}
|
||||
|
||||
auto record_batch =
|
||||
arrow::RecordBatch::Make(schema, range.row_count, std::move(arrays));
|
||||
auto validate_status = record_batch->Validate();
|
||||
if (!validate_status.ok()) {
|
||||
ThrowInfo(milvus::storage::ArrowStatusToErrorCode(validate_status),
|
||||
"invalid materialized json stats record batch: {}",
|
||||
validate_status.ToString());
|
||||
}
|
||||
|
||||
size_t materialized_bytes = sizeof(MaterializedChunk);
|
||||
for (const auto& array : record_batch->columns()) {
|
||||
materialized_bytes =
|
||||
SaturatingAdd(materialized_bytes, GetArrowArrayMemorySize(array));
|
||||
materialized_bytes = SaturatingAdd(
|
||||
materialized_bytes, sizeof(std::shared_ptr<arrow::Array>));
|
||||
}
|
||||
for (const auto& [key, postings] : bson_postings) {
|
||||
materialized_bytes = SaturatingAdd(
|
||||
materialized_bytes,
|
||||
sizeof(std::pair<const std::string, std::vector<int64_t>>) +
|
||||
3 * sizeof(void*));
|
||||
materialized_bytes =
|
||||
SaturatingAdd(materialized_bytes, key.capacity() + 1);
|
||||
materialized_bytes = SaturatingAdd(
|
||||
materialized_bytes,
|
||||
SaturatingMultiply(postings.capacity(), sizeof(int64_t)));
|
||||
}
|
||||
return MaterializedChunk{
|
||||
std::move(record_batch), std::move(bson_postings), materialized_bytes};
|
||||
}
|
||||
|
||||
void
|
||||
JsonKeyStats::BuildKeyStatsParallel(
|
||||
const std::vector<JsonStatsRowRange>& row_ranges, bool nullable) {
|
||||
if (row_ranges.empty()) {
|
||||
AssertInfo(num_rows_ == 0,
|
||||
"materialized 0 json stats rows, expected {}",
|
||||
num_rows_);
|
||||
return;
|
||||
}
|
||||
|
||||
AssertInfo(parquet_writer_ != nullptr,
|
||||
"json stats parquet writer must be initialized before "
|
||||
"materialize");
|
||||
AssertInfo(bson_inverted_index_ != nullptr,
|
||||
"json stats bson index must be initialized before materialize");
|
||||
auto schema = parquet_writer_->GetSchema();
|
||||
AssertInfo(schema != nullptr,
|
||||
"json stats parquet writer schema must be initialized before "
|
||||
"materialize");
|
||||
|
||||
auto* executor = milvus::futures::getJsonStatsBuildExecutor();
|
||||
const auto max_active_tasks = std::min(
|
||||
row_ranges.size(), std::max<size_t>(1, executor->numThreads()));
|
||||
auto& memory_budget = JsonStatsBuildMemoryBudget::GetInstance();
|
||||
|
||||
struct ActiveMaterializeTask {
|
||||
std::shared_ptr<std::atomic<size_t>> accounted_bytes;
|
||||
folly::Future<MaterializedChunk> future;
|
||||
};
|
||||
|
||||
std::vector<ActiveMaterializeTask> tasks;
|
||||
tasks.reserve(row_ranges.size());
|
||||
size_t next_submit = 0;
|
||||
size_t next_consume = 0;
|
||||
std::exception_ptr first_exception;
|
||||
|
||||
auto remember_exception = [&](std::exception_ptr exception) {
|
||||
if (first_exception == nullptr) {
|
||||
first_exception = std::move(exception);
|
||||
}
|
||||
};
|
||||
|
||||
auto submit_next = [&](bool block_for_budget) -> bool {
|
||||
size_t reservation_bytes = 0;
|
||||
JsonStatsRowRange range;
|
||||
try {
|
||||
range = row_ranges[next_submit];
|
||||
reservation_bytes =
|
||||
EstimateJsonStatsMaterializeReservationBytes(range, *schema);
|
||||
} catch (...) {
|
||||
remember_exception(std::current_exception());
|
||||
return false;
|
||||
}
|
||||
|
||||
std::shared_ptr<std::atomic<size_t>> accounted_bytes;
|
||||
try {
|
||||
accounted_bytes =
|
||||
std::make_shared<std::atomic<size_t>>(reservation_bytes);
|
||||
if (block_for_budget) {
|
||||
memory_budget.Acquire(reservation_bytes);
|
||||
} else if (!memory_budget.TryAcquire(reservation_bytes)) {
|
||||
return false;
|
||||
}
|
||||
} catch (...) {
|
||||
remember_exception(std::current_exception());
|
||||
return false;
|
||||
}
|
||||
|
||||
try {
|
||||
auto future = folly::via(
|
||||
executor,
|
||||
[this,
|
||||
range = std::move(range),
|
||||
nullable,
|
||||
schema,
|
||||
reservation_bytes,
|
||||
accounted_bytes,
|
||||
&memory_budget]() {
|
||||
auto result =
|
||||
MaterializeKeyStatsRange(range, nullable, schema);
|
||||
memory_budget.ReconcileReservation(reservation_bytes,
|
||||
result.memory_bytes);
|
||||
accounted_bytes->store(result.memory_bytes,
|
||||
std::memory_order_release);
|
||||
return result;
|
||||
});
|
||||
tasks.push_back(
|
||||
ActiveMaterializeTask{accounted_bytes, std::move(future)});
|
||||
++next_submit;
|
||||
} catch (...) {
|
||||
memory_budget.Release(
|
||||
accounted_bytes->load(std::memory_order_acquire));
|
||||
remember_exception(std::current_exception());
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
};
|
||||
|
||||
auto refill = [&]() {
|
||||
while (first_exception == nullptr && next_submit < row_ranges.size() &&
|
||||
tasks.size() - next_consume < max_active_tasks) {
|
||||
const bool block_for_budget = tasks.size() == next_consume;
|
||||
if (!submit_next(block_for_budget)) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
auto release_task_budget = [&](const ActiveMaterializeTask& task) {
|
||||
memory_budget.Release(
|
||||
task.accounted_bytes->load(std::memory_order_acquire));
|
||||
};
|
||||
|
||||
auto drain_submitted_tasks = [&]() {
|
||||
while (next_consume < tasks.size()) {
|
||||
auto& task = tasks[next_consume];
|
||||
try {
|
||||
std::move(task.future).get();
|
||||
} catch (...) {
|
||||
remember_exception(std::current_exception());
|
||||
}
|
||||
release_task_budget(task);
|
||||
++next_consume;
|
||||
}
|
||||
};
|
||||
|
||||
refill();
|
||||
|
||||
int64_t processed_rows = 0;
|
||||
while (next_consume < tasks.size()) {
|
||||
auto& task = tasks[next_consume];
|
||||
try {
|
||||
auto result = std::move(task.future).get();
|
||||
if (first_exception == nullptr) {
|
||||
const auto& range = row_ranges[next_consume];
|
||||
AssertInfo(range.global_begin == processed_rows,
|
||||
"json stats materialize result starts at global "
|
||||
"row {}, expected {}",
|
||||
range.global_begin,
|
||||
processed_rows);
|
||||
|
||||
auto status =
|
||||
parquet_writer_->AppendRecordBatch(result.record_batch);
|
||||
if (!status.ok()) {
|
||||
ThrowInfo(milvus::storage::ArrowStatusToErrorCode(status),
|
||||
"failed to append materialized json stats record "
|
||||
"batch: {}",
|
||||
status.ToString());
|
||||
}
|
||||
for (auto& [key, postings] : result.bson_postings) {
|
||||
bson_inverted_index_->AddRecords(key, std::move(postings));
|
||||
}
|
||||
|
||||
processed_rows += range.row_count;
|
||||
}
|
||||
} catch (...) {
|
||||
remember_exception(std::current_exception());
|
||||
}
|
||||
release_task_budget(task);
|
||||
++next_consume;
|
||||
|
||||
if (first_exception != nullptr) {
|
||||
drain_submitted_tasks();
|
||||
break;
|
||||
}
|
||||
|
||||
refill();
|
||||
}
|
||||
|
||||
// Consume every submitted future before propagating an error so no worker
|
||||
// can retain references to this JsonKeyStats or its FieldData inputs.
|
||||
if (first_exception != nullptr) {
|
||||
std::rethrow_exception(first_exception);
|
||||
}
|
||||
|
||||
AssertInfo(processed_rows == num_rows_,
|
||||
"materialized {} json stats rows, expected {}",
|
||||
processed_rows,
|
||||
num_rows_);
|
||||
}
|
||||
|
||||
std::string
|
||||
@@ -901,13 +1580,35 @@ JsonKeyStats::AddBucketName(const std::string& remote_prefix) {
|
||||
void
|
||||
JsonKeyStats::BuildWithFieldData(const std::vector<FieldDataPtr>& field_datas,
|
||||
bool nullable) {
|
||||
// collect key stats info and classify key type
|
||||
auto infos = CollectKeyInfo(field_datas, nullable);
|
||||
LOG_INFO("collect key infos: {} for segment {} for field {}",
|
||||
PrintKeyInfo(infos),
|
||||
segment_id_,
|
||||
field_id_);
|
||||
key_types_ = ClassifyJsonKeyLayoutType(infos);
|
||||
const auto row_ranges =
|
||||
CreateJsonStatsRowRanges(field_datas, kJsonStatsRowsPerRange);
|
||||
auto* executor = milvus::futures::getJsonStatsBuildExecutor();
|
||||
const auto executor_threads = std::max<size_t>(1, executor->numThreads());
|
||||
const auto effective_concurrency =
|
||||
std::min(row_ranges.size(), executor_threads);
|
||||
const bool use_parallel_pipeline = effective_concurrency > 1;
|
||||
LOG_INFO(
|
||||
"json stats build selects {} pipeline for segment {} field {} with "
|
||||
"{} row ranges, {} executor threads, and effective concurrency {}",
|
||||
use_parallel_pipeline ? "parallel" : "serial",
|
||||
segment_id_,
|
||||
field_id_,
|
||||
row_ranges.size(),
|
||||
executor_threads,
|
||||
effective_concurrency);
|
||||
|
||||
{
|
||||
// Release collected key stats after classification, before materializing
|
||||
// rows and accumulating BSON postings.
|
||||
auto infos = use_parallel_pipeline
|
||||
? CollectKeyInfo(row_ranges, nullable)
|
||||
: CollectKeyInfo(field_datas, nullable);
|
||||
LOG_INFO("collect key infos: {} for segment {} for field {}",
|
||||
PrintKeyInfo(infos),
|
||||
segment_id_,
|
||||
field_id_);
|
||||
key_types_ = ClassifyJsonKeyLayoutType(infos);
|
||||
}
|
||||
LOG_INFO("key types infos: {} for segment {} for field {}",
|
||||
PrintJsonKeyLayoutType(key_types_),
|
||||
segment_id_,
|
||||
@@ -931,7 +1632,20 @@ JsonKeyStats::BuildWithFieldData(const std::vector<FieldDataPtr>& field_datas,
|
||||
auto writer_context =
|
||||
ParquetWriterFactory::CreateContext(key_types_, remote_prefix);
|
||||
parquet_writer_->Init(std::move(writer_context));
|
||||
BuildKeyStats(field_datas, nullable);
|
||||
if (use_parallel_pipeline) {
|
||||
// Arrow slices retain the complete backing buffers. Keep every
|
||||
// materialized range within the writer batch limit so
|
||||
// AppendRecordBatch can hand it off without slicing; otherwise
|
||||
// downstream memory accounting would charge a small slice for the
|
||||
// full range and undersize row groups.
|
||||
const auto materialize_rows_per_range =
|
||||
std::min(kJsonStatsRowsPerRange, write_batch_size_);
|
||||
const auto materialize_row_ranges =
|
||||
CreateJsonStatsRowRanges(field_datas, materialize_rows_per_range);
|
||||
BuildKeyStatsParallel(materialize_row_ranges, nullable);
|
||||
} else {
|
||||
BuildKeyStats(row_ranges, nullable);
|
||||
}
|
||||
auto close_status = parquet_writer_->Close();
|
||||
if (!close_status.ok()) {
|
||||
ThrowInfo(milvus::storage::ArrowStatusToErrorCode(close_status),
|
||||
|
||||
@@ -29,6 +29,7 @@
|
||||
#include <mutex>
|
||||
#include <optional>
|
||||
#include <set>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
#include <string_view>
|
||||
#include <type_traits>
|
||||
@@ -71,20 +72,41 @@
|
||||
#include "storage/MemFileManagerImpl.h"
|
||||
|
||||
class CollectSingleJsonStatsInfoAccessor;
|
||||
class CollectKeyInfoAccessor;
|
||||
class BuildKeyStatsAccessor;
|
||||
// Forward declaration of test accessor in global namespace for friend declaration
|
||||
class TraverseJsonForBuildStatsAccessor;
|
||||
class JsonStatsProjectionTestAccessor;
|
||||
class JsonStatsScanTestAccessor;
|
||||
|
||||
namespace milvus::index {
|
||||
|
||||
struct JsonStatsFieldDataSlice {
|
||||
FieldDataPtr data;
|
||||
int64_t local_begin;
|
||||
int64_t row_count;
|
||||
};
|
||||
|
||||
struct JsonStatsRowRange {
|
||||
int64_t global_begin;
|
||||
int64_t row_count;
|
||||
std::vector<JsonStatsFieldDataSlice> slices;
|
||||
};
|
||||
|
||||
std::vector<JsonStatsRowRange>
|
||||
CreateJsonStatsRowRanges(const std::vector<FieldDataPtr>& field_datas,
|
||||
int64_t max_rows_per_range);
|
||||
|
||||
class JsonKeyStats : public ScalarIndex<std::string> {
|
||||
public:
|
||||
static constexpr int64_t kDefaultWriteBatchSize = 81920;
|
||||
|
||||
explicit JsonKeyStats(
|
||||
const storage::FileManagerContext& ctx,
|
||||
bool is_load,
|
||||
int64_t json_stats_max_shredding_columns = 1024,
|
||||
double json_stats_shredding_ratio_threshold = 0.3,
|
||||
int64_t json_stats_write_batch_size = 81920,
|
||||
int64_t json_stats_write_batch_size = kDefaultWriteBatchSize,
|
||||
uint32_t tantivy_index_version = TANTIVY_INDEX_LATEST_VERSION);
|
||||
|
||||
~JsonKeyStats() override;
|
||||
@@ -463,9 +485,14 @@ class JsonKeyStats : public ScalarIndex<std::string> {
|
||||
return ss.str();
|
||||
}
|
||||
|
||||
// Keep the original serial implementation for differential tests.
|
||||
std::map<JsonKey, KeyStatsInfo>
|
||||
CollectKeyInfo(const std::vector<FieldDataPtr>& field_datas, bool nullable);
|
||||
|
||||
std::map<JsonKey, KeyStatsInfo>
|
||||
CollectKeyInfo(const std::vector<JsonStatsRowRange>& row_ranges,
|
||||
bool nullable);
|
||||
|
||||
void
|
||||
TraverseJsonForStats(const char* json,
|
||||
jsmntok* tokens,
|
||||
@@ -496,8 +523,30 @@ class JsonKeyStats : public ScalarIndex<std::string> {
|
||||
std::map<JsonKey, JsonKeyLayoutType>
|
||||
ClassifyJsonKeyLayoutType(const std::map<JsonKey, KeyStatsInfo>& infos);
|
||||
|
||||
// Keep the original serial implementation for differential tests.
|
||||
void
|
||||
BuildKeyStats(const std::vector<FieldDataPtr>& field_datas, bool nullable);
|
||||
BuildKeyStats(const std::vector<JsonStatsRowRange>& row_ranges,
|
||||
bool nullable);
|
||||
|
||||
void
|
||||
BuildKeyStatsParallel(const std::vector<JsonStatsRowRange>& row_ranges,
|
||||
bool nullable);
|
||||
|
||||
static size_t
|
||||
EstimateJsonStatsMaterializeReservationBytes(const JsonStatsRowRange& range,
|
||||
const arrow::Schema& schema);
|
||||
|
||||
struct MaterializedChunk {
|
||||
std::shared_ptr<arrow::RecordBatch> record_batch;
|
||||
std::map<std::string, std::vector<int64_t>> bson_postings;
|
||||
size_t memory_bytes;
|
||||
};
|
||||
|
||||
MaterializedChunk
|
||||
MaterializeKeyStatsRange(
|
||||
const JsonStatsRowRange& range,
|
||||
bool nullable,
|
||||
const std::shared_ptr<arrow::Schema>& schema) const;
|
||||
|
||||
void
|
||||
BuildKeyStatsForRow(std::string_view json_str, uint32_t row_id);
|
||||
@@ -524,22 +573,26 @@ class JsonKeyStats : public ScalarIndex<std::string> {
|
||||
AddKeyStats(const std::vector<std::string>& path,
|
||||
JSONType type,
|
||||
const std::string& value,
|
||||
std::map<JsonKey, std::string>& values);
|
||||
std::map<JsonKey, std::string>& values) const;
|
||||
|
||||
bool
|
||||
ParseJsonForBuildStats(std::string_view json_str,
|
||||
std::map<JsonKey, std::string>& values) const;
|
||||
|
||||
void
|
||||
TraverseJsonForBuildStats(const char* json,
|
||||
jsmntok* tokens,
|
||||
int& index,
|
||||
std::vector<std::string>& path,
|
||||
std::map<JsonKey, std::string>& values);
|
||||
std::map<JsonKey, std::string>& values) const;
|
||||
|
||||
bool
|
||||
IsBoolean(const std::string& str) {
|
||||
IsBoolean(const std::string& str) const {
|
||||
return str == "true" || str == "false";
|
||||
}
|
||||
|
||||
bool
|
||||
IsInt64(const std::string& str) {
|
||||
IsInt64(const std::string& str) const {
|
||||
std::istringstream iss(str);
|
||||
int64_t num;
|
||||
iss >> num;
|
||||
@@ -548,32 +601,36 @@ class JsonKeyStats : public ScalarIndex<std::string> {
|
||||
}
|
||||
|
||||
bool
|
||||
IsFloat(const std::string& str) {
|
||||
IsFloat(const std::string& str) const {
|
||||
try {
|
||||
std::stof(str);
|
||||
return true;
|
||||
} catch (...) {
|
||||
} catch (const std::invalid_argument&) {
|
||||
return false;
|
||||
} catch (const std::out_of_range&) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
bool
|
||||
IsDouble(const std::string& str) {
|
||||
IsDouble(const std::string& str) const {
|
||||
try {
|
||||
std::stod(str);
|
||||
return true;
|
||||
} catch (...) {
|
||||
} catch (const std::invalid_argument&) {
|
||||
return false;
|
||||
} catch (const std::out_of_range&) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
bool
|
||||
IsNull(const std::string& str) {
|
||||
IsNull(const std::string& str) const {
|
||||
return str == "null";
|
||||
}
|
||||
|
||||
JSONType
|
||||
getType(const std::string& str) {
|
||||
getType(const std::string& str) const {
|
||||
if (IsBoolean(str)) {
|
||||
return JSONType::BOOL;
|
||||
// TODO: add int8, int16, int32 support
|
||||
@@ -684,6 +741,8 @@ class JsonKeyStats : public ScalarIndex<std::string> {
|
||||
// Friend accessor for unit tests to call private methods safely.
|
||||
friend class ::TraverseJsonForBuildStatsAccessor;
|
||||
friend class ::CollectSingleJsonStatsInfoAccessor;
|
||||
friend class ::CollectKeyInfoAccessor;
|
||||
friend class ::BuildKeyStatsAccessor;
|
||||
friend class ::JsonStatsProjectionTestAccessor;
|
||||
friend class ::JsonStatsScanTestAccessor;
|
||||
};
|
||||
|
||||
@@ -0,0 +1,148 @@
|
||||
// Licensed to the LF AI & Data foundation under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <condition_variable>
|
||||
#include <cstddef>
|
||||
#include <limits>
|
||||
#include <mutex>
|
||||
|
||||
#include "common/EasyAssert.h"
|
||||
|
||||
namespace milvus::index {
|
||||
|
||||
/// Process-wide byte budget for JSON stats ranges submitted for asynchronous
|
||||
/// materialization but not yet consumed. Capacity 0 means unlimited. An
|
||||
/// oversized range may run exclusively to guarantee progress.
|
||||
class JsonStatsBuildMemoryBudget {
|
||||
public:
|
||||
static JsonStatsBuildMemoryBudget&
|
||||
GetInstance() {
|
||||
static JsonStatsBuildMemoryBudget instance;
|
||||
return instance;
|
||||
}
|
||||
|
||||
/// Block until enough budget is available. The caller must not hold work
|
||||
/// whose completion is needed to release budget.
|
||||
void
|
||||
Acquire(size_t bytes) {
|
||||
std::unique_lock<std::mutex> lock(mu_);
|
||||
cv_.wait(lock, [this, bytes] { return CanAcquireLocked(bytes); });
|
||||
inflight_bytes_ += bytes;
|
||||
}
|
||||
|
||||
/// Try to claim budget without blocking. Refill loops use this while they
|
||||
/// still own inflight work.
|
||||
bool
|
||||
TryAcquire(size_t bytes) {
|
||||
std::lock_guard<std::mutex> lock(mu_);
|
||||
if (!CanAcquireLocked(bytes)) {
|
||||
return false;
|
||||
}
|
||||
inflight_bytes_ += bytes;
|
||||
return true;
|
||||
}
|
||||
|
||||
void
|
||||
Release(size_t bytes) {
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(mu_);
|
||||
AssertInfo(
|
||||
bytes <= inflight_bytes_,
|
||||
"JSON stats build memory budget over-release: release {}, "
|
||||
"inflight {}",
|
||||
bytes,
|
||||
inflight_bytes_);
|
||||
inflight_bytes_ -= bytes;
|
||||
}
|
||||
cv_.notify_all();
|
||||
}
|
||||
|
||||
/// Replace a task's pre-dispatch reservation with its measured result
|
||||
/// size. Growth is intentionally non-blocking: completed workers must not
|
||||
/// wait for consumers ordered behind other active tasks. A temporary
|
||||
/// overage prevents new acquisitions until consumers release enough bytes.
|
||||
void
|
||||
ReconcileReservation(size_t reserved_bytes, size_t actual_bytes) {
|
||||
bool released_bytes = false;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(mu_);
|
||||
AssertInfo(reserved_bytes <= inflight_bytes_,
|
||||
"JSON stats build memory budget reconcile exceeds "
|
||||
"inflight: "
|
||||
"reserved {}, inflight {}",
|
||||
reserved_bytes,
|
||||
inflight_bytes_);
|
||||
const auto other_inflight_bytes = inflight_bytes_ - reserved_bytes;
|
||||
AssertInfo(actual_bytes <= std::numeric_limits<size_t>::max() -
|
||||
other_inflight_bytes,
|
||||
"JSON stats build memory budget reconcile overflow: "
|
||||
"actual "
|
||||
"{}, other inflight {}",
|
||||
actual_bytes,
|
||||
other_inflight_bytes);
|
||||
inflight_bytes_ = other_inflight_bytes + actual_bytes;
|
||||
released_bytes = actual_bytes < reserved_bytes;
|
||||
}
|
||||
if (released_bytes) {
|
||||
cv_.notify_all();
|
||||
}
|
||||
}
|
||||
|
||||
size_t
|
||||
CapacityBytes() const {
|
||||
std::lock_guard<std::mutex> lock(mu_);
|
||||
return capacity_bytes_;
|
||||
}
|
||||
|
||||
size_t
|
||||
InflightBytes() const {
|
||||
std::lock_guard<std::mutex> lock(mu_);
|
||||
return inflight_bytes_;
|
||||
}
|
||||
|
||||
void
|
||||
SetCapacityBytes(size_t bytes) {
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(mu_);
|
||||
capacity_bytes_ = bytes;
|
||||
}
|
||||
cv_.notify_all();
|
||||
}
|
||||
|
||||
private:
|
||||
JsonStatsBuildMemoryBudget() = default;
|
||||
|
||||
bool
|
||||
CanAcquireLocked(size_t bytes) const {
|
||||
if (capacity_bytes_ == 0) {
|
||||
return true;
|
||||
}
|
||||
if (bytes > capacity_bytes_) {
|
||||
return inflight_bytes_ == 0;
|
||||
}
|
||||
return inflight_bytes_ <= capacity_bytes_ &&
|
||||
bytes <= capacity_bytes_ - inflight_bytes_;
|
||||
}
|
||||
|
||||
mutable std::mutex mu_;
|
||||
std::condition_variable cv_;
|
||||
size_t inflight_bytes_{0};
|
||||
size_t capacity_bytes_{0};
|
||||
};
|
||||
|
||||
} // namespace milvus::index
|
||||
@@ -76,6 +76,25 @@ BsonInvertedIndex::AddRecord(const std::string& key,
|
||||
EncodeInvertedIndexValue(row_id, offset));
|
||||
}
|
||||
|
||||
void
|
||||
BsonInvertedIndex::AddRecords(const std::string& key,
|
||||
std::vector<int64_t> records) {
|
||||
if (records.empty()) {
|
||||
return;
|
||||
}
|
||||
|
||||
auto& destination = inverted_index_map_[key];
|
||||
if (destination.empty()) {
|
||||
destination = std::move(records);
|
||||
return;
|
||||
}
|
||||
|
||||
// Let vector retain spare capacity from its geometric growth. Reserving
|
||||
// exactly the required size for every chunk would reallocate and copy the
|
||||
// entire accumulated postings list on nearly every merge.
|
||||
destination.insert(destination.end(), records.begin(), records.end());
|
||||
}
|
||||
|
||||
void
|
||||
BsonInvertedIndex::BuildIndex() {
|
||||
if (wrapper_ == nullptr) {
|
||||
|
||||
@@ -54,6 +54,9 @@ class BsonInvertedIndex {
|
||||
void
|
||||
AddRecord(const std::string& key, uint32_t row_id, uint32_t offset);
|
||||
|
||||
void
|
||||
AddRecords(const std::string& key, std::vector<int64_t> records);
|
||||
|
||||
void
|
||||
BuildIndex();
|
||||
|
||||
|
||||
@@ -18,7 +18,9 @@
|
||||
|
||||
#include <arrow/array/builder_binary.h>
|
||||
#include <arrow/array/builder_primitive.h>
|
||||
#include <algorithm>
|
||||
#include <exception>
|
||||
#include <new>
|
||||
#include <utility>
|
||||
|
||||
#include "arrow/array/builder_base.h"
|
||||
@@ -194,7 +196,7 @@ JsonStatsParquetWriter::AppendValue(const std::string& key,
|
||||
}
|
||||
|
||||
auto& builder = it->second;
|
||||
auto ast = AppendDataToBuilder(value, builder);
|
||||
auto ast = AppendJsonStatsValueToBuilder(value, builder);
|
||||
if (!ast.ok()) {
|
||||
ThrowInfo(milvus::storage::ArrowStatusToErrorCode(ast),
|
||||
"failed to append data to builder");
|
||||
@@ -213,7 +215,7 @@ JsonStatsParquetWriter::AppendRow(
|
||||
}
|
||||
|
||||
auto& builder = it->second;
|
||||
auto status = AppendDataToBuilder(value, builder);
|
||||
auto status = AppendJsonStatsValueToBuilder(value, builder);
|
||||
if (!status.ok()) {
|
||||
ThrowInfo(milvus::storage::ArrowStatusToErrorCode(status),
|
||||
"failed to append data to builder");
|
||||
@@ -224,7 +226,73 @@ JsonStatsParquetWriter::AppendRow(
|
||||
}
|
||||
|
||||
arrow::Status
|
||||
JsonStatsParquetWriter::AppendDataToBuilder(
|
||||
JsonStatsParquetWriter::AppendRecordBatch(
|
||||
const std::shared_ptr<arrow::RecordBatch>& batch) {
|
||||
if (batch == nullptr) {
|
||||
return arrow::Status::Invalid(
|
||||
"cannot append a null json stats record batch");
|
||||
}
|
||||
if (schema_ == nullptr || packed_writer_ == nullptr) {
|
||||
return arrow::Status::Invalid(
|
||||
"json stats parquet writer is not initialized");
|
||||
}
|
||||
if (schema_->num_fields() != batch->num_columns()) {
|
||||
return arrow::Status::Invalid(
|
||||
"json stats record batch column count does not match writer "
|
||||
"schema");
|
||||
}
|
||||
if (!schema_->Equals(*batch->schema(), true)) {
|
||||
return arrow::Status::Invalid(
|
||||
"json stats record batch schema does not match writer schema");
|
||||
}
|
||||
auto validate_status = batch->Validate();
|
||||
if (!validate_status.ok()) {
|
||||
return validate_status;
|
||||
}
|
||||
|
||||
if (batch->num_rows() == 0) {
|
||||
return arrow::Status::OK();
|
||||
}
|
||||
|
||||
// Preserve row order if a caller switches from the row-oriented API to
|
||||
// pre-materialized record batches. The parallel JSON-stats path does not
|
||||
// normally have pending builder rows.
|
||||
if (unflushed_row_count_ > 0) {
|
||||
auto result = WriteCurrentBatch();
|
||||
if (!result.ok()) {
|
||||
return result.status();
|
||||
}
|
||||
}
|
||||
|
||||
// MaterializeKeyStatsRange already produced complete Arrow arrays. Pass
|
||||
// them straight through so the writer does not copy every value into a
|
||||
// second set of builders merely to finish those builders back into arrays.
|
||||
// Keep batch_size_ as an upper bound for a single write; Slice shares the
|
||||
// original buffers and therefore remains zero-copy.
|
||||
const auto max_batch_rows = static_cast<int64_t>(batch_size_);
|
||||
int64_t offset = 0;
|
||||
while (offset < batch->num_rows()) {
|
||||
const auto rows_to_write =
|
||||
std::min(max_batch_rows, batch->num_rows() - offset);
|
||||
auto slice = offset == 0 && rows_to_write == batch->num_rows()
|
||||
? batch
|
||||
: batch->Slice(offset, rows_to_write);
|
||||
auto status = packed_writer_->Write(slice);
|
||||
if (!status.ok()) {
|
||||
return status;
|
||||
}
|
||||
offset += rows_to_write;
|
||||
}
|
||||
|
||||
// Account for the original arrays once. Multiple zero-copy slices share
|
||||
// these buffers, so accounting each slice would overstate their size.
|
||||
UpdatePathSizeMap(batch->columns());
|
||||
all_row_count_ += static_cast<size_t>(batch->num_rows());
|
||||
return arrow::Status::OK();
|
||||
}
|
||||
|
||||
arrow::Status
|
||||
AppendJsonStatsValueToBuilder(
|
||||
const std::string& value,
|
||||
const std::shared_ptr<arrow::ArrayBuilder>& builder) {
|
||||
auto type_id = builder->type()->id();
|
||||
@@ -296,6 +364,10 @@ JsonStatsParquetWriter::AppendDataToBuilder(
|
||||
type_id,
|
||||
value);
|
||||
}
|
||||
} catch (const std::bad_alloc&) {
|
||||
throw;
|
||||
} catch (const SegcoreError&) {
|
||||
throw;
|
||||
} catch (const std::exception& e) {
|
||||
ThrowInfo(ErrorCode::UnexpectedError,
|
||||
"failed to append data to builder: {}",
|
||||
|
||||
@@ -101,6 +101,11 @@ struct TableStatsInfo {
|
||||
std::map<JsonKey, JsonKeyLayoutType> column_map;
|
||||
};
|
||||
|
||||
arrow::Status
|
||||
AppendJsonStatsValueToBuilder(
|
||||
const std::string& value,
|
||||
const std::shared_ptr<arrow::ArrayBuilder>& builder);
|
||||
|
||||
class ColumnGroupingStrategy {
|
||||
public:
|
||||
virtual ~ColumnGroupingStrategy() = default;
|
||||
@@ -198,6 +203,9 @@ class JsonStatsParquetWriter {
|
||||
void
|
||||
AppendSharedRow(const uint8_t* data, size_t length);
|
||||
|
||||
arrow::Status
|
||||
AppendRecordBatch(const std::shared_ptr<arrow::RecordBatch>& batch);
|
||||
|
||||
arrow::Status
|
||||
Flush();
|
||||
|
||||
@@ -220,14 +228,15 @@ class JsonStatsParquetWriter {
|
||||
return total_size_;
|
||||
}
|
||||
|
||||
const std::shared_ptr<arrow::Schema>&
|
||||
GetSchema() const {
|
||||
return schema_;
|
||||
}
|
||||
|
||||
void
|
||||
UpdatePathSizeMap(const std::vector<std::shared_ptr<arrow::Array>>& arrays);
|
||||
|
||||
private:
|
||||
arrow::Status
|
||||
AppendDataToBuilder(const std::string& value,
|
||||
const std::shared_ptr<arrow::ArrayBuilder>& builder);
|
||||
|
||||
// init info
|
||||
std::shared_ptr<arrow::Schema> schema_{nullptr};
|
||||
|
||||
|
||||
@@ -151,7 +151,8 @@ CreateArrowBuilders(const std::map<JsonKey, JsonKeyLayoutType>& column_map) {
|
||||
for (const auto& [key, type] : column_map) {
|
||||
switch (type) {
|
||||
case JsonKeyLayoutType::SHARED:
|
||||
builders_map[key.ToColumnName()] = shared_builder;
|
||||
// Shared paths are appended together through the final BSON
|
||||
// builder, never looked up as independent Arrow columns.
|
||||
break;
|
||||
case JsonKeyLayoutType::TYPED:
|
||||
case JsonKeyLayoutType::TYPED_NOT_ALL:
|
||||
@@ -172,6 +173,21 @@ CreateArrowBuilders(const std::map<JsonKey, JsonKeyLayoutType>& column_map) {
|
||||
return {std::move(builders), std::move(builders_map)};
|
||||
}
|
||||
|
||||
std::pair<std::vector<std::shared_ptr<arrow::ArrayBuilder>>,
|
||||
std::map<std::string, std::shared_ptr<arrow::ArrayBuilder>>>
|
||||
CreateArrowBuildersForColumns(const std::set<JsonKey>& column_keys) {
|
||||
std::vector<std::shared_ptr<arrow::ArrayBuilder>> builders;
|
||||
builders.reserve(column_keys.size() + 1);
|
||||
std::map<std::string, std::shared_ptr<arrow::ArrayBuilder>> builders_map;
|
||||
for (const auto& key : column_keys) {
|
||||
auto builder = CreateArrowBuilder(key.type_);
|
||||
builders_map[key.ToColumnName()] = builder;
|
||||
builders.push_back(std::move(builder));
|
||||
}
|
||||
builders.push_back(CreateSharedArrowBuilder());
|
||||
return {std::move(builders), std::move(builders_map)};
|
||||
}
|
||||
|
||||
std::shared_ptr<arrow::Schema>
|
||||
CreateArrowSchema(const std::map<JsonKey, JsonKeyLayoutType>& column_map) {
|
||||
std::vector<std::shared_ptr<arrow::Field>> fields;
|
||||
|
||||
@@ -19,6 +19,7 @@
|
||||
#include <cstring>
|
||||
#include "common/FastMem.h"
|
||||
#include <map>
|
||||
#include <new>
|
||||
#include <optional>
|
||||
#include <set>
|
||||
#include <string>
|
||||
@@ -83,6 +84,8 @@ UnescapeJsonString(const std::string& escaped) {
|
||||
escaped);
|
||||
}
|
||||
return std::string(std::string_view(elem.get_string()));
|
||||
} catch (const std::bad_alloc&) {
|
||||
throw;
|
||||
} catch (const SegcoreError&) {
|
||||
// Already classified above (DataFormatBroken); SegcoreError derives from
|
||||
// std::runtime_error, so without this the generic handler below would
|
||||
@@ -398,6 +401,12 @@ std::pair<std::vector<std::shared_ptr<arrow::ArrayBuilder>>,
|
||||
std::map<std::string, std::shared_ptr<arrow::ArrayBuilder>>>
|
||||
CreateArrowBuilders(const std::map<JsonKey, JsonKeyLayoutType>& column_map);
|
||||
|
||||
// Build task-local Arrow columns without visiting the potentially much larger
|
||||
// set of keys stored in the shared BSON column. The shared builder is last.
|
||||
std::pair<std::vector<std::shared_ptr<arrow::ArrayBuilder>>,
|
||||
std::map<std::string, std::shared_ptr<arrow::ArrayBuilder>>>
|
||||
CreateArrowBuildersForColumns(const std::set<JsonKey>& column_keys);
|
||||
|
||||
std::shared_ptr<arrow::Schema>
|
||||
CreateArrowSchema(const std::map<JsonKey, JsonKeyLayoutType>& column_map);
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,8 +1,10 @@
|
||||
#include <arrow/type.h>
|
||||
#include <gmock/gmock.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <filesystem>
|
||||
#include <map>
|
||||
#include <memory>
|
||||
#include <new>
|
||||
#include <set>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
@@ -13,9 +15,73 @@
|
||||
#include "index/json_stats/parquet_writer.h"
|
||||
#include "index/json_stats/utils.h"
|
||||
#include "segcore/default_fs.h"
|
||||
#include "storage/StatusToErrorCode.h"
|
||||
#include "test_utils/Constants.h"
|
||||
|
||||
namespace milvus::index {
|
||||
namespace {
|
||||
|
||||
class MockMemoryPool : public arrow::ProxyMemoryPool {
|
||||
public:
|
||||
MockMemoryPool() : arrow::ProxyMemoryPool(arrow::default_memory_pool()) {
|
||||
}
|
||||
|
||||
using arrow::ProxyMemoryPool::Allocate;
|
||||
MOCK_METHOD(arrow::Status,
|
||||
Allocate,
|
||||
(int64_t, int64_t, uint8_t**),
|
||||
(override));
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
TEST(JsonStatsArrowBuilderTest, PreservesBadAlloc) {
|
||||
MockMemoryPool pool;
|
||||
auto builder = std::make_shared<arrow::StringBuilder>(&pool);
|
||||
EXPECT_CALL(pool, Allocate).WillOnce(testing::Throw(std::bad_alloc{}));
|
||||
|
||||
EXPECT_THROW(AppendJsonStatsValueToBuilder("value", builder),
|
||||
std::bad_alloc);
|
||||
}
|
||||
|
||||
TEST(JsonStatsArrowBuilderTest, PreservesTypedError) {
|
||||
MockMemoryPool pool;
|
||||
auto builder = std::make_shared<arrow::StringBuilder>(&pool);
|
||||
EXPECT_CALL(pool, Allocate)
|
||||
.WillOnce(testing::Throw(
|
||||
SegcoreError(MemAllocateFailed, "injected allocation failure")));
|
||||
|
||||
try {
|
||||
AppendJsonStatsValueToBuilder("value", builder);
|
||||
FAIL() << "expected an allocation failure";
|
||||
} catch (const SegcoreError& error) {
|
||||
EXPECT_EQ(error.get_error_code(), MemAllocateFailed);
|
||||
EXPECT_STREQ(error.what(), "injected allocation failure");
|
||||
}
|
||||
}
|
||||
|
||||
TEST(JsonStatsArrowBuilderTest, PreservesArrowOutOfMemory) {
|
||||
MockMemoryPool pool;
|
||||
auto builder = std::make_shared<arrow::StringBuilder>(&pool);
|
||||
EXPECT_CALL(pool, Allocate)
|
||||
.WillOnce(testing::Return(
|
||||
arrow::Status::OutOfMemory("injected allocation failure")));
|
||||
|
||||
auto status = AppendJsonStatsValueToBuilder("value", builder);
|
||||
EXPECT_TRUE(status.IsOutOfMemory());
|
||||
EXPECT_EQ(storage::ArrowStatusToErrorCode(status), MemAllocateFailed);
|
||||
}
|
||||
|
||||
TEST(JsonStatsArrowBuilderTest, PreservesUnsupportedType) {
|
||||
auto builder = std::make_shared<arrow::Date32Builder>();
|
||||
try {
|
||||
AppendJsonStatsValueToBuilder("1", builder);
|
||||
FAIL() << "expected an unsupported type error";
|
||||
} catch (const SegcoreError& error) {
|
||||
EXPECT_EQ(error.get_error_code(), Unsupported);
|
||||
}
|
||||
}
|
||||
|
||||
class ParquetWriterFactoryTest : public ::testing::Test {
|
||||
protected:
|
||||
void
|
||||
@@ -62,7 +128,7 @@ TEST_F(ParquetWriterFactoryTest, CreateContextBasicTest) {
|
||||
EXPECT_FALSE(context.builders.empty());
|
||||
EXPECT_FALSE(context.builders_map.empty());
|
||||
EXPECT_EQ(context.builders.size(), column_map_.size());
|
||||
EXPECT_EQ(context.builders_map.size(), column_map_.size());
|
||||
EXPECT_EQ(context.builders_map.size(), column_map_.size() - 1);
|
||||
|
||||
// Verify metadata
|
||||
EXPECT_TRUE(context.kv_metadata.empty());
|
||||
@@ -89,7 +155,7 @@ TEST_F(ParquetWriterFactoryTest, CreateContextWithSharedFields) {
|
||||
EXPECT_NE(context.schema, nullptr);
|
||||
EXPECT_EQ(context.schema->num_fields(), 2);
|
||||
|
||||
EXPECT_EQ(context.builders_map.size(), 3);
|
||||
EXPECT_EQ(context.builders_map.size(), 1);
|
||||
}
|
||||
|
||||
TEST_F(ParquetWriterFactoryTest, CreateContextWithColumnGroups) {
|
||||
@@ -142,4 +208,63 @@ TEST_F(ParquetWriterFactoryTest, CloseReturnsStatusAndIsIdempotent) {
|
||||
std::filesystem::remove_all(path_prefix);
|
||||
}
|
||||
|
||||
TEST_F(ParquetWriterFactoryTest,
|
||||
AppendRecordBatchDoesNotRematerializeIntoWriterBuilders) {
|
||||
auto fs = milvus::segcore::GetDefaultArrowFileSystem();
|
||||
auto path_prefix = std::filesystem::path(TestLocalPath) /
|
||||
"json_stats_writer_direct_record_batch";
|
||||
std::filesystem::remove_all(path_prefix);
|
||||
ASSERT_TRUE(std::filesystem::create_directories(path_prefix));
|
||||
|
||||
std::map<JsonKey, JsonKeyLayoutType> column_map = {
|
||||
{JsonKey("/int", JSONType::INT64), JsonKeyLayoutType::TYPED},
|
||||
{JsonKey("/shared", JSONType::STRING), JsonKeyLayoutType::SHARED},
|
||||
};
|
||||
auto context =
|
||||
ParquetWriterFactory::CreateContext(column_map, path_prefix.string());
|
||||
auto schema = context.schema;
|
||||
auto writer_builders = context.builders;
|
||||
|
||||
milvus_storage::StorageConfig storage_config;
|
||||
JsonStatsParquetWriter writer(fs, storage_config, 16 * 1024 * 1024, 1024);
|
||||
writer.Init(std::move(context));
|
||||
|
||||
writer.AppendValue(JsonKey("/int", JSONType::INT64).ToColumnName(), "5");
|
||||
writer.AppendSharedRow(nullptr, 0);
|
||||
writer.AddCurrentRow();
|
||||
for (const auto& builder : writer_builders) {
|
||||
ASSERT_EQ(builder->length(), 1);
|
||||
}
|
||||
|
||||
auto input_builders = CreateArrowBuilders(column_map).first;
|
||||
auto int_builder =
|
||||
std::static_pointer_cast<arrow::Int64Builder>(input_builders.front());
|
||||
auto shared_builder =
|
||||
std::static_pointer_cast<arrow::BinaryBuilder>(input_builders.back());
|
||||
for (int64_t value : {int64_t{10}, int64_t{20}, int64_t{30}}) {
|
||||
ASSERT_TRUE(int_builder->Append(value).ok());
|
||||
ASSERT_TRUE(shared_builder->AppendNull().ok());
|
||||
}
|
||||
|
||||
std::vector<std::shared_ptr<arrow::Array>> arrays;
|
||||
arrays.reserve(input_builders.size());
|
||||
for (auto& builder : input_builders) {
|
||||
std::shared_ptr<arrow::Array> array;
|
||||
ASSERT_TRUE(builder->Finish(&array).ok());
|
||||
arrays.push_back(std::move(array));
|
||||
}
|
||||
auto batch = arrow::RecordBatch::Make(schema, 3, std::move(arrays));
|
||||
|
||||
auto status = writer.AppendRecordBatch(batch);
|
||||
ASSERT_TRUE(status.ok()) << status.ToString();
|
||||
for (const auto& builder : writer_builders) {
|
||||
EXPECT_EQ(builder->length(), 0);
|
||||
}
|
||||
EXPECT_FALSE(writer.GetPathsToSize().empty());
|
||||
|
||||
status = writer.Close();
|
||||
ASSERT_TRUE(status.ok()) << status.ToString();
|
||||
std::filesystem::remove_all(path_prefix);
|
||||
}
|
||||
|
||||
} // namespace milvus::index
|
||||
|
||||
@@ -8,10 +8,14 @@
|
||||
// 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 <cstddef>
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
#include <map>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <string_view>
|
||||
#include <tuple>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
@@ -19,6 +23,8 @@
|
||||
#include "common/Json.h"
|
||||
#include "common/jsmn.h"
|
||||
#include "common/protobuf_utils.h"
|
||||
#include "futures/Executor.h"
|
||||
#include "futures/future_c.h"
|
||||
#include "gtest/gtest.h"
|
||||
#include "index/json_stats/JsonKeyStats.h"
|
||||
#include "index/json_stats/utils.h"
|
||||
@@ -57,6 +63,25 @@ class CollectSingleJsonStatsInfoAccessor {
|
||||
}
|
||||
};
|
||||
|
||||
// Friend accessor declared in JsonKeyStats to invoke private method for UT
|
||||
class CollectKeyInfoAccessor {
|
||||
public:
|
||||
static std::map<JsonKey, milvus::index::KeyStatsInfo>
|
||||
CallLegacy(JsonKeyStats& stats,
|
||||
const std::vector<milvus::FieldDataPtr>& field_datas,
|
||||
bool nullable) {
|
||||
return stats.CollectKeyInfo(field_datas, nullable);
|
||||
}
|
||||
|
||||
static std::map<JsonKey, milvus::index::KeyStatsInfo>
|
||||
CallParallel(
|
||||
JsonKeyStats& stats,
|
||||
const std::vector<milvus::index::JsonStatsRowRange>& row_ranges,
|
||||
bool nullable) {
|
||||
return stats.CollectKeyInfo(row_ranges, nullable);
|
||||
}
|
||||
};
|
||||
|
||||
namespace {
|
||||
|
||||
// Helper to tokenize JSON using jsmn
|
||||
@@ -81,8 +106,523 @@ Tokenize(const char* json) {
|
||||
return tokens;
|
||||
}
|
||||
|
||||
milvus::FieldDataPtr
|
||||
CreateJsonFieldData(int64_t row_count) {
|
||||
auto field_data = milvus::storage::CreateFieldData(
|
||||
milvus::DataType::JSON, milvus::DataType::NONE, false);
|
||||
if (row_count == 0) {
|
||||
return field_data;
|
||||
}
|
||||
|
||||
std::vector<milvus::Json> rows;
|
||||
rows.reserve(row_count);
|
||||
const std::string json = R"({"value":1})";
|
||||
for (int64_t i = 0; i < row_count; ++i) {
|
||||
rows.emplace_back(simdjson::padded_string(json));
|
||||
}
|
||||
field_data->FillFieldData(rows.data(), rows.size());
|
||||
return field_data;
|
||||
}
|
||||
|
||||
milvus::FieldDataPtr
|
||||
CreateJsonFieldData(const std::vector<std::string>& json_rows) {
|
||||
auto field_data = milvus::storage::CreateFieldData(
|
||||
milvus::DataType::JSON, milvus::DataType::NONE, false);
|
||||
if (json_rows.empty()) {
|
||||
return field_data;
|
||||
}
|
||||
|
||||
std::vector<milvus::Json> rows;
|
||||
rows.reserve(json_rows.size());
|
||||
for (const auto& json : json_rows) {
|
||||
rows.emplace_back(simdjson::padded_string(json));
|
||||
}
|
||||
field_data->FillFieldData(rows.data(), rows.size());
|
||||
return field_data;
|
||||
}
|
||||
|
||||
milvus::FieldDataPtr
|
||||
CreateNullableJsonFieldData(const std::vector<std::string>& json_rows,
|
||||
const std::vector<bool>& valid_rows) {
|
||||
if (json_rows.size() != valid_rows.size()) {
|
||||
ADD_FAILURE() << "json rows and validity rows must have the same size";
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
auto field_data = milvus::storage::CreateFieldData(
|
||||
milvus::DataType::JSON, milvus::DataType::NONE, true);
|
||||
if (json_rows.empty()) {
|
||||
return field_data;
|
||||
}
|
||||
|
||||
std::vector<milvus::Json> rows;
|
||||
rows.reserve(json_rows.size());
|
||||
for (const auto& json : json_rows) {
|
||||
rows.emplace_back(simdjson::padded_string(json));
|
||||
}
|
||||
|
||||
std::vector<uint8_t> valid_bitmap((valid_rows.size() + 7) / 8, 0);
|
||||
for (size_t i = 0; i < valid_rows.size(); ++i) {
|
||||
if (valid_rows[i]) {
|
||||
valid_bitmap[i / 8] |= static_cast<uint8_t>(1U << (i % 8));
|
||||
}
|
||||
}
|
||||
field_data->FillFieldData(rows.data(), valid_bitmap.data(), rows.size(), 0);
|
||||
return field_data;
|
||||
}
|
||||
|
||||
std::unique_ptr<JsonKeyStats>
|
||||
CreateJsonKeyStatsForCollectTest() {
|
||||
milvus::storage::FieldDataMeta field_meta{1, 2, 3, 100, {}};
|
||||
milvus::storage::IndexMeta index_meta{3, 100, 1, 1};
|
||||
milvus::storage::StorageConfig storage_config;
|
||||
storage_config.storage_type = "local";
|
||||
storage_config.root_path = TestLocalPath;
|
||||
auto cm = milvus::storage::CreateChunkManager(storage_config);
|
||||
auto fs = milvus::storage::InitArrowFileSystem(storage_config);
|
||||
milvus::storage::FileManagerContext ctx(field_meta, index_meta, cm, fs);
|
||||
return std::make_unique<JsonKeyStats>(ctx, true);
|
||||
}
|
||||
|
||||
void
|
||||
ExpectKeyInfoMapsEqual(
|
||||
const std::map<JsonKey, milvus::index::KeyStatsInfo>& expected,
|
||||
const std::map<JsonKey, milvus::index::KeyStatsInfo>& actual) {
|
||||
ASSERT_EQ(actual.size(), expected.size());
|
||||
for (const auto& [key, expected_info] : expected) {
|
||||
auto actual_it = actual.find(key);
|
||||
ASSERT_NE(actual_it, actual.end()) << "missing key " << key.ToString();
|
||||
EXPECT_EQ(actual_it->second.hit_row_num_, expected_info.hit_row_num_)
|
||||
<< key.ToString();
|
||||
EXPECT_EQ(std::memcmp(actual_it->second.min_value_,
|
||||
expected_info.min_value_,
|
||||
sizeof(expected_info.min_value_)),
|
||||
0)
|
||||
<< key.ToString();
|
||||
EXPECT_EQ(std::memcmp(actual_it->second.max_value_,
|
||||
expected_info.max_value_,
|
||||
sizeof(expected_info.max_value_)),
|
||||
0)
|
||||
<< key.ToString();
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
ExpectHitCount(const std::map<JsonKey, milvus::index::KeyStatsInfo>& infos,
|
||||
const JsonKey& key,
|
||||
int32_t expected_hit_count) {
|
||||
auto it = infos.find(key);
|
||||
ASSERT_NE(it, infos.end()) << "missing key " << key.ToString();
|
||||
EXPECT_EQ(it->second.hit_row_num_, expected_hit_count) << key.ToString();
|
||||
}
|
||||
|
||||
std::map<JsonKey, milvus::index::KeyStatsInfo>
|
||||
ExpectParallelCollectMatchesLegacy(
|
||||
const std::vector<milvus::FieldDataPtr>& field_datas,
|
||||
bool nullable,
|
||||
const std::vector<int64_t>& rows_per_range_values,
|
||||
const std::vector<int>& worker_counts) {
|
||||
int64_t expected_row_count = 0;
|
||||
for (const auto& data : field_datas) {
|
||||
expected_row_count += data->get_num_rows();
|
||||
}
|
||||
|
||||
auto stats = CreateJsonKeyStatsForCollectTest();
|
||||
auto legacy_infos =
|
||||
CollectKeyInfoAccessor::CallLegacy(*stats, field_datas, nullable);
|
||||
EXPECT_EQ(stats->Count(), expected_row_count);
|
||||
|
||||
for (auto rows_per_range : rows_per_range_values) {
|
||||
auto ranges = milvus::index::CreateJsonStatsRowRanges(field_datas,
|
||||
rows_per_range);
|
||||
for (auto worker_count : worker_counts) {
|
||||
SCOPED_TRACE(::testing::Message()
|
||||
<< "rows_per_range=" << rows_per_range
|
||||
<< ", worker_count=" << worker_count
|
||||
<< ", nullable=" << nullable);
|
||||
executor_set_json_stats_build_thread_num(worker_count);
|
||||
auto parallel_infos =
|
||||
CollectKeyInfoAccessor::CallParallel(*stats, ranges, nullable);
|
||||
EXPECT_EQ(stats->Count(), expected_row_count);
|
||||
ExpectKeyInfoMapsEqual(legacy_infos, parallel_infos);
|
||||
}
|
||||
}
|
||||
return legacy_infos;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
class JsonStatsCollectKeyInfoTest : public ::testing::Test {
|
||||
protected:
|
||||
void
|
||||
SetUp() override {
|
||||
original_thread_num_ =
|
||||
milvus::futures::getJsonStatsBuildExecutor()->numThreads();
|
||||
}
|
||||
|
||||
void
|
||||
TearDown() override {
|
||||
executor_set_json_stats_build_thread_num(
|
||||
static_cast<int>(original_thread_num_));
|
||||
}
|
||||
|
||||
size_t original_thread_num_;
|
||||
};
|
||||
|
||||
TEST(JsonStatsRowRangeTest, CoalescesAcrossFieldDataBoundaries) {
|
||||
auto first = CreateJsonFieldData(2);
|
||||
auto second = CreateJsonFieldData(5);
|
||||
auto third = CreateJsonFieldData(1);
|
||||
|
||||
auto ranges =
|
||||
milvus::index::CreateJsonStatsRowRanges({first, second, third}, 3);
|
||||
|
||||
ASSERT_EQ(ranges.size(), 3);
|
||||
|
||||
EXPECT_EQ(ranges[0].global_begin, 0);
|
||||
EXPECT_EQ(ranges[0].row_count, 3);
|
||||
ASSERT_EQ(ranges[0].slices.size(), 2);
|
||||
EXPECT_EQ(ranges[0].slices[0].data.get(), first.get());
|
||||
EXPECT_EQ(ranges[0].slices[0].local_begin, 0);
|
||||
EXPECT_EQ(ranges[0].slices[0].row_count, 2);
|
||||
EXPECT_EQ(ranges[0].slices[1].data.get(), second.get());
|
||||
EXPECT_EQ(ranges[0].slices[1].local_begin, 0);
|
||||
EXPECT_EQ(ranges[0].slices[1].row_count, 1);
|
||||
|
||||
EXPECT_EQ(ranges[1].global_begin, 3);
|
||||
EXPECT_EQ(ranges[1].row_count, 3);
|
||||
ASSERT_EQ(ranges[1].slices.size(), 1);
|
||||
EXPECT_EQ(ranges[1].slices[0].data.get(), second.get());
|
||||
EXPECT_EQ(ranges[1].slices[0].local_begin, 1);
|
||||
EXPECT_EQ(ranges[1].slices[0].row_count, 3);
|
||||
|
||||
EXPECT_EQ(ranges[2].global_begin, 6);
|
||||
EXPECT_EQ(ranges[2].row_count, 2);
|
||||
ASSERT_EQ(ranges[2].slices.size(), 2);
|
||||
EXPECT_EQ(ranges[2].slices[0].data.get(), second.get());
|
||||
EXPECT_EQ(ranges[2].slices[0].local_begin, 4);
|
||||
EXPECT_EQ(ranges[2].slices[0].row_count, 1);
|
||||
EXPECT_EQ(ranges[2].slices[1].data.get(), third.get());
|
||||
EXPECT_EQ(ranges[2].slices[1].local_begin, 0);
|
||||
EXPECT_EQ(ranges[2].slices[1].row_count, 1);
|
||||
}
|
||||
|
||||
TEST(JsonStatsRowRangeTest, MatchesLegacyTraversalForBoundarySizes) {
|
||||
std::vector<milvus::FieldDataPtr> field_datas = {
|
||||
CreateJsonFieldData(0),
|
||||
CreateJsonFieldData(2),
|
||||
CreateJsonFieldData(0),
|
||||
CreateJsonFieldData(5),
|
||||
CreateJsonFieldData(1),
|
||||
CreateJsonFieldData(0),
|
||||
};
|
||||
|
||||
using RowPosition =
|
||||
std::tuple<const milvus::FieldDataBase*, int64_t, int64_t>;
|
||||
std::vector<RowPosition> legacy_rows;
|
||||
int64_t global_row = 0;
|
||||
for (const auto& data : field_datas) {
|
||||
for (int64_t local_row = 0; local_row < data->get_num_rows();
|
||||
++local_row) {
|
||||
legacy_rows.emplace_back(data.get(), local_row, global_row++);
|
||||
}
|
||||
}
|
||||
|
||||
for (int64_t max_rows_per_range : {1, 2, 3, 4, 8, 16}) {
|
||||
auto ranges = milvus::index::CreateJsonStatsRowRanges(
|
||||
field_datas, max_rows_per_range);
|
||||
std::vector<RowPosition> ranged_rows;
|
||||
int64_t expected_global_begin = 0;
|
||||
|
||||
for (const auto& range : ranges) {
|
||||
EXPECT_EQ(range.global_begin, expected_global_begin);
|
||||
EXPECT_GT(range.row_count, 0);
|
||||
EXPECT_LE(range.row_count, max_rows_per_range);
|
||||
|
||||
int64_t rows_in_range = 0;
|
||||
for (const auto& slice : range.slices) {
|
||||
const auto slice_global_begin =
|
||||
range.global_begin + rows_in_range;
|
||||
for (int64_t offset = 0; offset < slice.row_count; ++offset) {
|
||||
ranged_rows.emplace_back(slice.data.get(),
|
||||
slice.local_begin + offset,
|
||||
slice_global_begin + offset);
|
||||
}
|
||||
rows_in_range += slice.row_count;
|
||||
}
|
||||
EXPECT_EQ(rows_in_range, range.row_count);
|
||||
expected_global_begin += range.row_count;
|
||||
}
|
||||
|
||||
EXPECT_EQ(ranged_rows, legacy_rows);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(JsonStatsRowRangeTest, HandlesEmptyInputAndRejectsZeroRangeSize) {
|
||||
EXPECT_TRUE(milvus::index::CreateJsonStatsRowRanges({}, 4).empty());
|
||||
EXPECT_TRUE(
|
||||
milvus::index::CreateJsonStatsRowRanges({CreateJsonFieldData(0)}, 4)
|
||||
.empty());
|
||||
EXPECT_ANY_THROW(
|
||||
milvus::index::CreateJsonStatsRowRanges({CreateJsonFieldData(1)}, 0));
|
||||
}
|
||||
|
||||
TEST_F(JsonStatsCollectKeyInfoTest,
|
||||
ParallelMatchesLegacyAcrossFieldDataBoundaries) {
|
||||
auto first = CreateJsonFieldData(std::vector<std::string>{
|
||||
R"({"a":1,"b":"x","nested":{"flag":true}})",
|
||||
R"({"a":2,"b":"y"})",
|
||||
R"({"a":"mixed","items":[1,2]})",
|
||||
});
|
||||
auto second = CreateJsonFieldData(std::vector<std::string>{
|
||||
R"({})",
|
||||
R"({"nested":{"flag":false},"null_value":null})",
|
||||
R"({"a":3,"b":"z"})",
|
||||
R"([])",
|
||||
"",
|
||||
});
|
||||
std::vector<milvus::FieldDataPtr> field_datas = {first, second};
|
||||
auto ranges = milvus::index::CreateJsonStatsRowRanges(field_datas, 2);
|
||||
ASSERT_EQ(ranges.size(), 4);
|
||||
ASSERT_EQ(ranges[1].slices.size(), 2);
|
||||
|
||||
auto stats = CreateJsonKeyStatsForCollectTest();
|
||||
auto legacy_infos =
|
||||
CollectKeyInfoAccessor::CallLegacy(*stats, field_datas, false);
|
||||
EXPECT_EQ(stats->Count(), 8);
|
||||
|
||||
executor_set_json_stats_build_thread_num(1);
|
||||
auto single_worker_infos =
|
||||
CollectKeyInfoAccessor::CallParallel(*stats, ranges, false);
|
||||
EXPECT_EQ(stats->Count(), 8);
|
||||
ExpectKeyInfoMapsEqual(legacy_infos, single_worker_infos);
|
||||
|
||||
executor_set_json_stats_build_thread_num(4);
|
||||
auto parallel_infos =
|
||||
CollectKeyInfoAccessor::CallParallel(*stats, ranges, false);
|
||||
EXPECT_EQ(stats->Count(), 8);
|
||||
ExpectKeyInfoMapsEqual(legacy_infos, parallel_infos);
|
||||
|
||||
ExpectHitCount(parallel_infos, JsonKey{"/a", JSONType::INT64}, 3);
|
||||
ExpectHitCount(parallel_infos, JsonKey{"/a", JSONType::STRING}, 1);
|
||||
ExpectHitCount(parallel_infos, JsonKey{"/nested", JSONType::OBJECT}, 2);
|
||||
ExpectHitCount(parallel_infos, JsonKey{"/nested/flag", JSONType::BOOL}, 2);
|
||||
ExpectHitCount(parallel_infos, JsonKey{"/null_value", JSONType::NONE}, 1);
|
||||
ExpectHitCount(parallel_infos, JsonKey{"", JSONType::ARRAY}, 1);
|
||||
}
|
||||
|
||||
TEST_F(JsonStatsCollectKeyInfoTest,
|
||||
ParallelPreservesNonNulTerminatedJsonViewLength) {
|
||||
const std::string json = R"({"a":1,"b":"x"})";
|
||||
std::string buffer = json;
|
||||
buffer.append(simdjson::SIMDJSON_PADDING, '\xAB');
|
||||
|
||||
auto field_data = std::dynamic_pointer_cast<milvus::FieldDataJsonImpl>(
|
||||
milvus::storage::CreateFieldData(
|
||||
milvus::DataType::JSON, milvus::DataType::NONE, false));
|
||||
ASSERT_NE(field_data, nullptr);
|
||||
field_data->add_json_data({milvus::Json(buffer.data(), json.size())});
|
||||
|
||||
auto ranges = milvus::index::CreateJsonStatsRowRanges({field_data}, 1);
|
||||
auto stats = CreateJsonKeyStatsForCollectTest();
|
||||
executor_set_json_stats_build_thread_num(2);
|
||||
auto infos = CollectKeyInfoAccessor::CallParallel(*stats, ranges, false);
|
||||
|
||||
EXPECT_EQ(stats->Count(), 1);
|
||||
ExpectHitCount(infos, JsonKey{"/a", JSONType::INT64}, 1);
|
||||
ExpectHitCount(infos, JsonKey{"/b", JSONType::STRING}, 1);
|
||||
}
|
||||
|
||||
TEST_F(JsonStatsCollectKeyInfoTest, ParallelMatchesLegacyForEmptyInputs) {
|
||||
auto empty_infos =
|
||||
ExpectParallelCollectMatchesLegacy({}, false, {1, 16 * 1024}, {1, 4});
|
||||
EXPECT_TRUE(empty_infos.empty());
|
||||
|
||||
std::vector<milvus::FieldDataPtr> empty_field_datas = {
|
||||
CreateJsonFieldData(0),
|
||||
CreateJsonFieldData(0),
|
||||
CreateJsonFieldData(0),
|
||||
};
|
||||
auto empty_field_data_infos = ExpectParallelCollectMatchesLegacy(
|
||||
empty_field_datas, false, {1, 2, 16 * 1024}, {1, 4});
|
||||
EXPECT_TRUE(empty_field_data_infos.empty());
|
||||
}
|
||||
|
||||
TEST_F(JsonStatsCollectKeyInfoTest,
|
||||
ParallelMatchesLegacyForAllJsonTypesAndDeepPaths) {
|
||||
auto first = CreateJsonFieldData(std::vector<std::string>{
|
||||
"42",
|
||||
"-42",
|
||||
"0",
|
||||
"1.5",
|
||||
"6.02e23",
|
||||
"true",
|
||||
"false",
|
||||
"null",
|
||||
"\"root string\"",
|
||||
"[]",
|
||||
R"([1,{"array_child":"is_not_traversed"}])",
|
||||
"{}",
|
||||
});
|
||||
auto second = CreateJsonFieldData(std::vector<std::string>{
|
||||
R"({"l1":{"l2":{"l3":{"int":1,"double":2.5,"bool":true,"null":null,"string":"x","array":[1,{"ignored":2}]}}},"a/b":{"~key":"escaped"}})",
|
||||
R"({"same":1})",
|
||||
R"({"same":"1"})",
|
||||
R"({"same":false})",
|
||||
R"({"same":null})",
|
||||
"",
|
||||
});
|
||||
std::vector<milvus::FieldDataPtr> field_datas = {
|
||||
first,
|
||||
CreateJsonFieldData(0),
|
||||
second,
|
||||
};
|
||||
|
||||
auto infos = ExpectParallelCollectMatchesLegacy(
|
||||
field_datas, false, {1, 2, 5, 16, 64}, {1, 2, 4});
|
||||
|
||||
ExpectHitCount(infos, JsonKey{"", JSONType::INT64}, 3);
|
||||
ExpectHitCount(infos, JsonKey{"", JSONType::DOUBLE}, 2);
|
||||
ExpectHitCount(infos, JsonKey{"", JSONType::BOOL}, 2);
|
||||
ExpectHitCount(infos, JsonKey{"", JSONType::NONE}, 1);
|
||||
ExpectHitCount(infos, JsonKey{"", JSONType::STRING}, 1);
|
||||
ExpectHitCount(infos, JsonKey{"", JSONType::ARRAY}, 2);
|
||||
ExpectHitCount(infos, JsonKey{"/l1/l2/l3/int", JSONType::INT64}, 1);
|
||||
ExpectHitCount(infos, JsonKey{"/a~1b/~0key", JSONType::STRING}, 1);
|
||||
ExpectHitCount(infos, JsonKey{"/same", JSONType::INT64}, 1);
|
||||
ExpectHitCount(infos, JsonKey{"/same", JSONType::STRING}, 1);
|
||||
ExpectHitCount(infos, JsonKey{"/same", JSONType::BOOL}, 1);
|
||||
ExpectHitCount(infos, JsonKey{"/same", JSONType::NONE}, 1);
|
||||
}
|
||||
|
||||
TEST_F(JsonStatsCollectKeyInfoTest, ParallelMatchesLegacyForNullableRows) {
|
||||
auto first = CreateNullableJsonFieldData(
|
||||
{
|
||||
R"({"kept":1})",
|
||||
R"({"broken":)",
|
||||
R"({"kept":2,"nested":{"flag":true}})",
|
||||
"not-json",
|
||||
},
|
||||
{true, false, true, false});
|
||||
auto second = CreateNullableJsonFieldData(
|
||||
{
|
||||
R"({"ignored":1})",
|
||||
R"({"kept":"two"})",
|
||||
R"({"nested":{"flag":false}})",
|
||||
},
|
||||
{false, true, true});
|
||||
auto non_nullable =
|
||||
CreateJsonFieldData(std::vector<std::string>{R"({"always":true})"});
|
||||
std::vector<milvus::FieldDataPtr> field_datas = {
|
||||
first,
|
||||
CreateNullableJsonFieldData({}, {}),
|
||||
second,
|
||||
non_nullable,
|
||||
};
|
||||
|
||||
auto infos_with_field_nullable = ExpectParallelCollectMatchesLegacy(
|
||||
field_datas, false, {1, 2, 3, 8}, {1, 4});
|
||||
auto infos_with_explicit_nullable = ExpectParallelCollectMatchesLegacy(
|
||||
field_datas, true, {1, 2, 3, 8}, {1, 4});
|
||||
ExpectKeyInfoMapsEqual(infos_with_field_nullable,
|
||||
infos_with_explicit_nullable);
|
||||
|
||||
ExpectHitCount(
|
||||
infos_with_field_nullable, JsonKey{"/kept", JSONType::INT64}, 2);
|
||||
ExpectHitCount(
|
||||
infos_with_field_nullable, JsonKey{"/kept", JSONType::STRING}, 1);
|
||||
ExpectHitCount(
|
||||
infos_with_field_nullable, JsonKey{"/nested/flag", JSONType::BOOL}, 2);
|
||||
ExpectHitCount(
|
||||
infos_with_field_nullable, JsonKey{"/always", JSONType::BOOL}, 1);
|
||||
EXPECT_EQ(
|
||||
infos_with_field_nullable.count(JsonKey{"/ignored", JSONType::INT64}),
|
||||
0);
|
||||
}
|
||||
|
||||
TEST_F(JsonStatsCollectKeyInfoTest, ParallelMatchesLegacyAcrossManyRangeSizes) {
|
||||
std::vector<milvus::FieldDataPtr> field_datas = {
|
||||
CreateJsonFieldData(0),
|
||||
CreateJsonFieldData(std::vector<std::string>{R"({"value":1})"}),
|
||||
CreateJsonFieldData(0),
|
||||
CreateJsonFieldData(std::vector<std::string>{
|
||||
R"({"value":2})",
|
||||
R"({"value":3})",
|
||||
R"({"value":"three"})",
|
||||
R"({"nested":{"value":4}})",
|
||||
R"({"array":[1,2,3]})",
|
||||
R"({"flag":true})",
|
||||
R"({"null_value":null})",
|
||||
}),
|
||||
CreateJsonFieldData(0),
|
||||
CreateJsonFieldData(std::vector<std::string>{
|
||||
R"({"value":5})",
|
||||
R"({"value":6})",
|
||||
R"({"value":7})",
|
||||
}),
|
||||
};
|
||||
|
||||
auto infos = ExpectParallelCollectMatchesLegacy(
|
||||
field_datas, false, {1, 2, 3, 4, 5, 8, 11, 32}, {1, 2, 4});
|
||||
ExpectHitCount(infos, JsonKey{"/value", JSONType::INT64}, 6);
|
||||
ExpectHitCount(infos, JsonKey{"/value", JSONType::STRING}, 1);
|
||||
}
|
||||
|
||||
TEST_F(JsonStatsCollectKeyInfoTest,
|
||||
ParallelMatchesLegacyForLargeMultiRangeInput) {
|
||||
constexpr int64_t kRowsPerRange = 16 * 1024;
|
||||
std::vector<milvus::FieldDataPtr> field_datas = {
|
||||
CreateJsonFieldData(kRowsPerRange - 1),
|
||||
CreateJsonFieldData(2),
|
||||
CreateJsonFieldData(kRowsPerRange),
|
||||
};
|
||||
|
||||
auto infos = ExpectParallelCollectMatchesLegacy(
|
||||
field_datas, false, {kRowsPerRange}, {1, 4});
|
||||
ExpectHitCount(
|
||||
infos, JsonKey{"/value", JSONType::INT64}, 2 * kRowsPerRange + 1);
|
||||
}
|
||||
|
||||
TEST_F(JsonStatsCollectKeyInfoTest, LegacyAndParallelRethrowSameParseError) {
|
||||
auto invalid =
|
||||
CreateJsonFieldData(std::vector<std::string>{R"({"broken":)"});
|
||||
auto valid = CreateJsonFieldData(
|
||||
std::vector<std::string>{R"({"value":1})", R"({"value":2})"});
|
||||
std::vector<milvus::FieldDataPtr> field_datas = {invalid, valid};
|
||||
auto ranges = milvus::index::CreateJsonStatsRowRanges(field_datas, 1);
|
||||
|
||||
milvus::ErrorCode legacy_error_code = milvus::UnexpectedError;
|
||||
std::string legacy_error_message;
|
||||
auto legacy_stats = CreateJsonKeyStatsForCollectTest();
|
||||
try {
|
||||
CollectKeyInfoAccessor::CallLegacy(*legacy_stats, field_datas, false);
|
||||
FAIL() << "expected invalid JSON to fail in legacy collect";
|
||||
} catch (const milvus::SegcoreError& error) {
|
||||
legacy_error_code = error.get_error_code();
|
||||
legacy_error_message = error.what();
|
||||
EXPECT_EQ(legacy_error_code, milvus::DataFormatBroken);
|
||||
EXPECT_NE(legacy_error_message.find("Failed to parse Json"),
|
||||
std::string::npos);
|
||||
}
|
||||
|
||||
for (int worker_count : {1, 4}) {
|
||||
SCOPED_TRACE(::testing::Message() << "worker_count=" << worker_count);
|
||||
executor_set_json_stats_build_thread_num(worker_count);
|
||||
auto parallel_stats = CreateJsonKeyStatsForCollectTest();
|
||||
try {
|
||||
CollectKeyInfoAccessor::CallParallel(
|
||||
*parallel_stats, ranges, false);
|
||||
FAIL() << "expected invalid JSON to fail in parallel collect";
|
||||
} catch (const milvus::SegcoreError& error) {
|
||||
EXPECT_EQ(error.get_error_code(), legacy_error_code);
|
||||
EXPECT_EQ(std::string(error.what()), legacy_error_message);
|
||||
}
|
||||
EXPECT_EQ(
|
||||
milvus::futures::getJsonStatsBuildExecutor()->getTaskQueueSize(),
|
||||
0);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(TraverseJsonForBuildStatsTest,
|
||||
HandlesPrimitivesArraysNestedAndEmptyObject) {
|
||||
const char* json = R"(
|
||||
|
||||
@@ -7,6 +7,7 @@
|
||||
#include <map>
|
||||
#include <memory>
|
||||
#include <optional>
|
||||
#include <set>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <utility>
|
||||
@@ -50,13 +51,55 @@ TEST_F(UtilsTest, CreateArrowBuildersTest) {
|
||||
|
||||
auto [builders, builders_map] = CreateArrowBuilders(column_map);
|
||||
EXPECT_EQ(builders.size(), column_map.size());
|
||||
EXPECT_EQ(builders_map.size(), column_map.size());
|
||||
EXPECT_EQ(builders_map.size(), column_map.size() - 1);
|
||||
EXPECT_EQ(builders_map.count(
|
||||
JsonKey("shared_key", JSONType::STRING).ToColumnName()),
|
||||
0);
|
||||
|
||||
auto schema = CreateArrowSchema(column_map);
|
||||
EXPECT_NE(schema, nullptr);
|
||||
EXPECT_EQ(schema->num_fields(), column_map.size());
|
||||
}
|
||||
|
||||
TEST_F(UtilsTest, SharedKeysDoNotExpandTaskLocalBuilders) {
|
||||
std::map<JsonKey, JsonKeyLayoutType> column_map = {
|
||||
{JsonKey("/typed", JSONType::INT64), JsonKeyLayoutType::TYPED},
|
||||
{JsonKey("/dynamic", JSONType::STRING), JsonKeyLayoutType::DYNAMIC},
|
||||
};
|
||||
std::set<JsonKey> column_keys;
|
||||
for (const auto& entry : column_map) {
|
||||
column_keys.insert(entry.first);
|
||||
}
|
||||
for (size_t i = 0; i < 4096; ++i) {
|
||||
column_map.emplace(
|
||||
JsonKey("/shared/" + std::to_string(i), JSONType::STRING),
|
||||
JsonKeyLayoutType::SHARED);
|
||||
}
|
||||
|
||||
auto [writer_builders, writer_map] = CreateArrowBuilders(column_map);
|
||||
auto [task_builders, task_map] = CreateArrowBuildersForColumns(column_keys);
|
||||
auto schema = CreateArrowSchema(column_map);
|
||||
|
||||
ASSERT_NE(schema, nullptr);
|
||||
ASSERT_EQ(schema->num_fields(), 3);
|
||||
ASSERT_EQ(writer_builders.size(), 3);
|
||||
ASSERT_EQ(task_builders.size(), 3);
|
||||
EXPECT_EQ(writer_map.size(), column_keys.size());
|
||||
EXPECT_EQ(task_map.size(), column_keys.size());
|
||||
EXPECT_EQ(
|
||||
writer_map.count(JsonKey("/shared/0", JSONType::STRING).ToColumnName()),
|
||||
0);
|
||||
size_t column_index = 0;
|
||||
for (const auto& key : column_keys) {
|
||||
const auto name = key.ToColumnName();
|
||||
EXPECT_EQ(schema->field(column_index)->name(), name);
|
||||
EXPECT_EQ(writer_map.at(name), writer_builders[column_index]);
|
||||
EXPECT_EQ(task_map.at(name), task_builders[column_index]);
|
||||
++column_index;
|
||||
}
|
||||
EXPECT_EQ(task_builders.back()->type()->id(), arrow::Type::BINARY);
|
||||
}
|
||||
|
||||
TEST_F(UtilsTest, CreateParquetKVMetadataTest) {
|
||||
std::map<JsonKey, JsonKeyLayoutType> column_map = {
|
||||
{JsonKey("int_key", JSONType::INT64), JsonKeyLayoutType::TYPED},
|
||||
|
||||
@@ -70,6 +70,14 @@ func InitSegcore(nodeID int64) error {
|
||||
|
||||
cCPUNum := C.int(hardware.GetCPUNum())
|
||||
C.InitCpuNum(cCPUNum)
|
||||
if err := initcore.UpdateJSONStatsBuildExecutorPoolSize(
|
||||
initcore.ResolveJSONStatsBuildExecutorPoolSize(paramtable.Get())); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := initcore.UpdateJSONStatsBuildMaxInflightBytes(
|
||||
initcore.ResolveJSONStatsBuildMaxInflightBytes(paramtable.Get())); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
cKnowhereThreadPoolSize := C.uint32_t(hardware.GetCPUNum() * paramtable.DefaultKnowhereThreadPoolNumRatioInBuild)
|
||||
if paramtable.GetRole() == typeutil.StandaloneRole {
|
||||
@@ -121,6 +129,8 @@ func InitSegcore(nodeID int64) error {
|
||||
initcore.RegisterArrowIOThreadPoolWatchers(paramtable.Get(), "datanode")
|
||||
initcore.RegisterArrowReaderConfigWatchers(paramtable.Get(), "datanode")
|
||||
initcore.RegisterLoonReaderConfigWatchers(paramtable.Get(), "datanode")
|
||||
initcore.RegisterJSONStatsBuildExecutorWatcher(paramtable.Get(), "datanode")
|
||||
initcore.RegisterJSONStatsBuildMemoryBudgetWatcher(paramtable.Get(), "datanode")
|
||||
|
||||
// init paramtable change callback for core related config
|
||||
initcore.SetupCoreConfigChangelCallback()
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
// Licensed to the LF AI & Data foundation under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package initcore
|
||||
|
||||
/*
|
||||
#cgo pkg-config: milvus_core
|
||||
|
||||
#include "futures/future_c.h"
|
||||
*/
|
||||
import "C"
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/milvus-io/milvus/pkg/v3/config"
|
||||
"github.com/milvus-io/milvus/pkg/v3/mlog"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/hardware"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
|
||||
)
|
||||
|
||||
const defaultJSONStatsBuildMaxWorkers = 8
|
||||
|
||||
// ResolveJSONStatsBuildExecutorPoolSize returns the process-wide hard worker
|
||||
// limit after applying the safe automatic default and the cgroup-aware CPU cap.
|
||||
func ResolveJSONStatsBuildExecutorPoolSize(pt *paramtable.ComponentParam) int {
|
||||
cpuNum := hardware.GetCPUNum()
|
||||
if cpuNum < 1 {
|
||||
cpuNum = 1
|
||||
}
|
||||
|
||||
maxWorkers := pt.DataNodeCfg.JSONStatsBuildMaxWorkers.GetAsInt()
|
||||
if maxWorkers <= 0 {
|
||||
maxWorkers = defaultJSONStatsBuildMaxWorkers
|
||||
}
|
||||
if maxWorkers > cpuNum {
|
||||
maxWorkers = cpuNum
|
||||
}
|
||||
return maxWorkers
|
||||
}
|
||||
|
||||
// UpdateJSONStatsBuildExecutorPoolSize resizes the C++ Folly executor.
|
||||
func UpdateJSONStatsBuildExecutorPoolSize(threads int) error {
|
||||
if threads < 1 {
|
||||
threads = 1
|
||||
}
|
||||
status := C.executor_set_json_stats_build_thread_num(C.int(threads))
|
||||
return HandleCStatus(&status, "failed to resize json stats build executor")
|
||||
}
|
||||
|
||||
// ResolveJSONStatsBuildMaxInflightBytes returns the process-wide materialized
|
||||
// chunk budget. Zero disables the byte limit.
|
||||
func ResolveJSONStatsBuildMaxInflightBytes(pt *paramtable.ComponentParam) int64 {
|
||||
maxInflightBytes := pt.DataNodeCfg.JSONStatsBuildMaxInflightBytes.GetAsInt64()
|
||||
if maxInflightBytes < 0 {
|
||||
return paramtable.DefaultJSONStatsBuildMaxInflightBytes
|
||||
}
|
||||
return maxInflightBytes
|
||||
}
|
||||
|
||||
// UpdateJSONStatsBuildMaxInflightBytes updates the C++ shared byte budget.
|
||||
func UpdateJSONStatsBuildMaxInflightBytes(maxInflightBytes int64) error {
|
||||
if maxInflightBytes < 0 {
|
||||
maxInflightBytes = paramtable.DefaultJSONStatsBuildMaxInflightBytes
|
||||
}
|
||||
status := C.executor_set_json_stats_build_max_inflight_bytes(C.int64_t(maxInflightBytes))
|
||||
return HandleCStatus(&status, "failed to update json stats build memory budget")
|
||||
}
|
||||
|
||||
// RegisterJSONStatsBuildExecutorWatcher applies maxWorkers updates without
|
||||
// interrupting tasks already running in the executor.
|
||||
func RegisterJSONStatsBuildExecutorWatcher(pt *paramtable.ComponentParam, source string) {
|
||||
key := pt.DataNodeCfg.JSONStatsBuildMaxWorkers.Key
|
||||
handler := func(evt *config.Event) {
|
||||
if !evt.HasUpdated {
|
||||
return
|
||||
}
|
||||
threads := ResolveJSONStatsBuildExecutorPoolSize(pt)
|
||||
if err := UpdateJSONStatsBuildExecutorPoolSize(threads); err != nil {
|
||||
mlog.Warn(context.TODO(), "failed to resize json stats build executor",
|
||||
mlog.String("source", source),
|
||||
mlog.String("trigger", key),
|
||||
mlog.Int("threads", threads),
|
||||
mlog.Err(err))
|
||||
return
|
||||
}
|
||||
mlog.Info(context.TODO(), "json stats build executor resized",
|
||||
mlog.String("source", source),
|
||||
mlog.String("trigger", key),
|
||||
mlog.Int("threads", threads))
|
||||
}
|
||||
pt.Watch(key, config.NewHandler("json-stats-build-executor."+source+"."+key, handler))
|
||||
}
|
||||
|
||||
// RegisterJSONStatsBuildMemoryBudgetWatcher applies maxInflightBytes updates.
|
||||
// Shrinking the budget does not interrupt active chunks; it only blocks later
|
||||
// dispatch until current reservations fall below the new capacity.
|
||||
func RegisterJSONStatsBuildMemoryBudgetWatcher(pt *paramtable.ComponentParam, source string) {
|
||||
key := pt.DataNodeCfg.JSONStatsBuildMaxInflightBytes.Key
|
||||
handler := func(evt *config.Event) {
|
||||
if !evt.HasUpdated {
|
||||
return
|
||||
}
|
||||
maxInflightBytes := ResolveJSONStatsBuildMaxInflightBytes(pt)
|
||||
if err := UpdateJSONStatsBuildMaxInflightBytes(maxInflightBytes); err != nil {
|
||||
mlog.Warn(context.TODO(), "failed to resize json stats build memory budget",
|
||||
mlog.String("source", source),
|
||||
mlog.String("trigger", key),
|
||||
mlog.Int64("maxInflightBytes", maxInflightBytes),
|
||||
mlog.Err(err))
|
||||
return
|
||||
}
|
||||
mlog.Info(context.TODO(), "json stats build memory budget resized",
|
||||
mlog.String("source", source),
|
||||
mlog.String("trigger", key),
|
||||
mlog.Int64("maxInflightBytes", maxInflightBytes))
|
||||
}
|
||||
pt.Watch(key, config.NewHandler("json-stats-build-memory-budget."+source+"."+key, handler))
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
// Licensed to the LF AI & Data foundation under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package initcore
|
||||
|
||||
import (
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/milvus-io/milvus/pkg/v3/config"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/hardware"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
|
||||
)
|
||||
|
||||
func TestJSONStatsBuildExecutorConfig(t *testing.T) {
|
||||
paramtable.Init()
|
||||
pt := paramtable.Get()
|
||||
key := pt.DataNodeCfg.JSONStatsBuildMaxWorkers.Key
|
||||
require.NoError(t, pt.Reset(key))
|
||||
defer pt.Reset(key)
|
||||
|
||||
cpuNum := hardware.GetCPUNum()
|
||||
if cpuNum < 1 {
|
||||
cpuNum = 1
|
||||
}
|
||||
|
||||
assert.Equal(t, min(defaultJSONStatsBuildMaxWorkers, cpuNum),
|
||||
ResolveJSONStatsBuildExecutorPoolSize(pt))
|
||||
|
||||
require.NoError(t, pt.Save(key, "2"))
|
||||
assert.Equal(t, min(2, cpuNum), ResolveJSONStatsBuildExecutorPoolSize(pt))
|
||||
|
||||
require.NoError(t, pt.Save(key, "0"))
|
||||
assert.Equal(t, min(defaultJSONStatsBuildMaxWorkers, cpuNum),
|
||||
ResolveJSONStatsBuildExecutorPoolSize(pt))
|
||||
require.NoError(t, UpdateJSONStatsBuildExecutorPoolSize(1))
|
||||
|
||||
RegisterJSONStatsBuildExecutorWatcher(pt, "test")
|
||||
var watcherFires atomic.Int32
|
||||
sentinel := config.NewHandler("json-stats-build-executor-sentinel", func(*config.Event) {
|
||||
watcherFires.Add(1)
|
||||
})
|
||||
pt.Watch(key, sentinel)
|
||||
defer pt.Unwatch(key, sentinel)
|
||||
|
||||
require.NoError(t, pt.Save(key, "1"))
|
||||
assert.Positive(t, watcherFires.Load())
|
||||
}
|
||||
|
||||
func TestJSONStatsBuildMemoryBudgetConfig(t *testing.T) {
|
||||
paramtable.Init()
|
||||
pt := paramtable.Get()
|
||||
key := pt.DataNodeCfg.JSONStatsBuildMaxInflightBytes.Key
|
||||
require.NoError(t, pt.Reset(key))
|
||||
defer pt.Reset(key)
|
||||
|
||||
assert.Equal(t, int64(paramtable.DefaultJSONStatsBuildMaxInflightBytes),
|
||||
ResolveJSONStatsBuildMaxInflightBytes(pt))
|
||||
|
||||
require.NoError(t, pt.Save(key, "67108864"))
|
||||
assert.Equal(t, int64(67108864), ResolveJSONStatsBuildMaxInflightBytes(pt))
|
||||
|
||||
require.NoError(t, pt.Save(key, "0"))
|
||||
assert.Zero(t, ResolveJSONStatsBuildMaxInflightBytes(pt))
|
||||
|
||||
require.NoError(t, pt.Save(key, "-1"))
|
||||
assert.Equal(t, int64(paramtable.DefaultJSONStatsBuildMaxInflightBytes),
|
||||
ResolveJSONStatsBuildMaxInflightBytes(pt))
|
||||
|
||||
require.NoError(t, UpdateJSONStatsBuildMaxInflightBytes(0))
|
||||
require.NoError(t, UpdateJSONStatsBuildMaxInflightBytes(128*1024*1024))
|
||||
|
||||
RegisterJSONStatsBuildMemoryBudgetWatcher(pt, "test")
|
||||
var watcherFires atomic.Int32
|
||||
sentinel := config.NewHandler("json-stats-build-memory-budget-sentinel", func(*config.Event) {
|
||||
watcherFires.Add(1)
|
||||
})
|
||||
pt.Watch(key, sentinel)
|
||||
defer pt.Unwatch(key, sentinel)
|
||||
|
||||
require.NoError(t, pt.Save(key, "33554432"))
|
||||
assert.Positive(t, watcherFires.Load())
|
||||
}
|
||||
@@ -46,8 +46,9 @@ const (
|
||||
DefaultIndexSliceSize = 16
|
||||
// Load admission defaults apply only when async loading is enabled and the
|
||||
// corresponding parameter is absent. Explicit values, including 0, win.
|
||||
DefaultLoadTransientBudgetBytes = 2 * 1024 * 1024 * 1024
|
||||
DefaultLoadAdmissionSlotsPerCPU = 2
|
||||
DefaultLoadTransientBudgetBytes = 2 * 1024 * 1024 * 1024
|
||||
DefaultLoadAdmissionSlotsPerCPU = 2
|
||||
DefaultJSONStatsBuildMaxInflightBytes = 512 * 1024 * 1024
|
||||
// DefaultStorageV2AsyncLoadReadWindowSizeBytes is the historical-key
|
||||
// default for the Storage V3 async read-window threshold.
|
||||
DefaultStorageV2AsyncLoadReadWindowSizeBytes = 16 * 1024 * 1024
|
||||
@@ -8088,13 +8089,21 @@ re-ingesting. A job that timed out before applying carries 0 and left the collec
|
||||
}
|
||||
p.JSONStatsShreddingRatioThreshold.Init(base.mgr)
|
||||
|
||||
const defaultJSONStatsWriteBatchSize = "81920"
|
||||
p.JSONStatsWriteBatchSize = ParamItem{
|
||||
Key: "dataCoord.jsonShreddingWriteBatchSize",
|
||||
Version: "2.6.5",
|
||||
DefaultValue: "81920",
|
||||
Doc: "the batch size to write",
|
||||
DefaultValue: defaultJSONStatsWriteBatchSize,
|
||||
Doc: "the batch size to write; values must be positive, otherwise the default is used",
|
||||
FallbackKeys: []string{"dataCoord.jsonStatsWriteBatchSize"},
|
||||
Export: true,
|
||||
Formatter: func(value string) string {
|
||||
batchSize, err := strconv.ParseInt(value, 10, 64)
|
||||
if err != nil || batchSize <= 0 {
|
||||
return defaultJSONStatsWriteBatchSize
|
||||
}
|
||||
return value
|
||||
},
|
||||
}
|
||||
p.JSONStatsWriteBatchSize.Init(base.mgr)
|
||||
}
|
||||
@@ -8187,8 +8196,10 @@ type dataNodeConfig struct {
|
||||
DeltalogFormat ParamItem `refreshable:"false"`
|
||||
|
||||
// index services config
|
||||
BuildParallel ParamItem `refreshable:"false"`
|
||||
MaxVecIndexBuildConcurrency ParamItem `refreshable:"true"`
|
||||
BuildParallel ParamItem `refreshable:"false"`
|
||||
MaxVecIndexBuildConcurrency ParamItem `refreshable:"true"`
|
||||
JSONStatsBuildMaxWorkers ParamItem `refreshable:"true"`
|
||||
JSONStatsBuildMaxInflightBytes ParamItem `refreshable:"true"`
|
||||
|
||||
WorkerSlotUnit ParamItem `refreshable:"true"`
|
||||
StandaloneSlotRatio ParamItem `refreshable:"false"`
|
||||
@@ -8748,6 +8759,37 @@ writeRetryInitialInterval, otherwise the effective cap is raised to twice the in
|
||||
}
|
||||
p.MaxVecIndexBuildConcurrency.Init(base.mgr)
|
||||
|
||||
p.JSONStatsBuildMaxWorkers = ParamItem{
|
||||
Key: "dataNode.jsonStats.parallelBuild.maxWorkers",
|
||||
Version: "3.0.1",
|
||||
DefaultValue: "8",
|
||||
Doc: "Maximum number of process-wide JSON Stats Build workers. The effective value is capped by the cgroup-aware CPU count. Values less than or equal to 0 use the safe automatic limit min(CPU, 8).",
|
||||
Export: true,
|
||||
}
|
||||
p.JSONStatsBuildMaxWorkers.Init(base.mgr)
|
||||
|
||||
p.JSONStatsBuildMaxInflightBytes = ParamItem{
|
||||
Key: "dataNode.jsonStats.parallelBuild.maxInflightBytes",
|
||||
Version: "3.0.1",
|
||||
DefaultValue: strconv.Itoa(DefaultJSONStatsBuildMaxInflightBytes),
|
||||
Doc: `Process-wide transient memory budget in bytes for materialized JSON Stats Build chunks. ` +
|
||||
`The budget is shared by concurrent builds. Lower values reduce peak memory at the cost of ` +
|
||||
`parallelism. A chunk larger than the budget is allowed to run exclusively to guarantee ` +
|
||||
`progress. Set to 0 to disable the limit.`,
|
||||
Export: true,
|
||||
Formatter: func(v string) string {
|
||||
value, err := strconv.ParseInt(v, 10, 64)
|
||||
if err != nil || value < 0 {
|
||||
mlog.Warn(context.TODO(), "dataNode.jsonStats.parallelBuild.maxInflightBytes must be non-negative, using default",
|
||||
mlog.String("configured", v),
|
||||
mlog.Int64("default", DefaultJSONStatsBuildMaxInflightBytes))
|
||||
return strconv.Itoa(DefaultJSONStatsBuildMaxInflightBytes)
|
||||
}
|
||||
return v
|
||||
},
|
||||
}
|
||||
p.JSONStatsBuildMaxInflightBytes.Init(base.mgr)
|
||||
|
||||
p.WorkerSlotUnit = ParamItem{
|
||||
Key: "dataNode.workerSlotUnit",
|
||||
Version: "2.5.7",
|
||||
|
||||
@@ -129,6 +129,33 @@ func TestComponentParam_DataCoordSnapshotExportCopyConcurrency(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestComponentParam_DataCoordJSONStatsWriteBatchSize(t *testing.T) {
|
||||
Init()
|
||||
params := Get()
|
||||
item := ¶ms.DataCoordCfg.JSONStatsWriteBatchSize
|
||||
legacyKey := item.FallbackKeys[0]
|
||||
params.Reset(item.Key)
|
||||
params.Reset(legacyKey)
|
||||
t.Cleanup(func() {
|
||||
params.Reset(item.Key)
|
||||
params.Reset(legacyKey)
|
||||
})
|
||||
|
||||
assert.EqualValues(t, 81920, item.GetAsInt64())
|
||||
|
||||
params.Save(item.Key, "4096")
|
||||
assert.EqualValues(t, 4096, item.GetAsInt64())
|
||||
|
||||
for _, invalid := range []string{"0", "-1", "invalid", "9223372036854775808"} {
|
||||
params.Save(item.Key, invalid)
|
||||
assert.EqualValues(t, 81920, item.GetAsInt64(), invalid)
|
||||
}
|
||||
|
||||
params.Reset(item.Key)
|
||||
params.Save(legacyKey, "-1")
|
||||
assert.EqualValues(t, 81920, item.GetAsInt64())
|
||||
}
|
||||
|
||||
func TestMembershipFilterConfig(t *testing.T) {
|
||||
base := NewBaseTable(SkipRemote(true))
|
||||
params := proxyConfig{}
|
||||
@@ -1260,6 +1287,22 @@ func TestComponentParam(t *testing.T) {
|
||||
assert.Equal(t, 10, Params.MaxCompactionConcurrency.GetAsInt())
|
||||
|
||||
assert.Equal(t, 4, Params.MaxVecIndexBuildConcurrency.GetAsInt())
|
||||
assert.Equal(t, 8, Params.JSONStatsBuildMaxWorkers.GetAsInt())
|
||||
assert.Equal(t, int64(DefaultJSONStatsBuildMaxInflightBytes),
|
||||
Params.JSONStatsBuildMaxInflightBytes.GetAsInt64())
|
||||
params.Save(Params.JSONStatsBuildMaxInflightBytes.Key, "-1")
|
||||
assert.Equal(t, int64(DefaultJSONStatsBuildMaxInflightBytes),
|
||||
Params.JSONStatsBuildMaxInflightBytes.GetAsInt64())
|
||||
params.Save(Params.JSONStatsBuildMaxInflightBytes.Key, "invalid")
|
||||
assert.Equal(t, int64(DefaultJSONStatsBuildMaxInflightBytes),
|
||||
Params.JSONStatsBuildMaxInflightBytes.GetAsInt64())
|
||||
params.Save(Params.JSONStatsBuildMaxInflightBytes.Key, "0")
|
||||
assert.Equal(t, int64(0),
|
||||
Params.JSONStatsBuildMaxInflightBytes.GetAsInt64())
|
||||
params.Save(Params.JSONStatsBuildMaxInflightBytes.Key, "67108864")
|
||||
assert.Equal(t, int64(67108864),
|
||||
Params.JSONStatsBuildMaxInflightBytes.GetAsInt64())
|
||||
params.Reset(Params.JSONStatsBuildMaxInflightBytes.Key)
|
||||
|
||||
// clustering compaction
|
||||
params.Save("datanode.clusteringCompaction.memoryBufferRatio", "0.1")
|
||||
|
||||
@@ -85,6 +85,7 @@ arrow internal/core/src/common/Schema.cpp
|
||||
arrow internal/core/src/common/Schema.h
|
||||
arrow internal/core/src/common/Types.h
|
||||
arrow internal/core/src/index/json_stats/JsonKeyStats.cpp
|
||||
arrow internal/core/src/index/json_stats/JsonKeyStats.h
|
||||
arrow internal/core/src/index/json_stats/parquet_writer.cpp
|
||||
arrow internal/core/src/index/json_stats/parquet_writer.h
|
||||
arrow internal/core/src/index/json_stats/utils.cpp
|
||||
|
||||
Reference in New Issue
Block a user