diff --git a/configs/milvus.yaml b/configs/milvus.yaml index c47125c0bb..c430ab72ee 100644 --- a/configs/milvus.yaml +++ b/configs/milvus.yaml @@ -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: diff --git a/internal/core/src/futures/Executor.cpp b/internal/core/src/futures/Executor.cpp index 833c2b3dc1..732ca3c1bf 100644 --- a/internal/core/src/futures/Executor.cpp +++ b/internal/core/src/futures/Executor.cpp @@ -10,7 +10,9 @@ // or implied. See the License for the specific language governing permissions and limitations under the License #include +#include #include +#include #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,9 +45,21 @@ getLoadCPUExecutor() { return &executor; } +folly::CPUThreadPoolExecutor* +getJsonStatsBuildExecutor() { + auto max_thread_num = static_cast(std::max(1, milvus::CPU_NUM)); + static folly::CPUThreadPoolExecutor executor( + std::pair{max_thread_num, kJsonStatsBuildMinThreads}, + folly::CPUThreadPoolExecutor::makeDefaultPriorityQueue( + kJsonStatsBuildNumPriority), + std::make_shared( + "MILVUS_JSON_STATS_BUILD_")); + return &executor; +} + folly::CPUThreadPoolExecutor* getGlobalCPUExecutor() { return getSearchCPUExecutor(); } -}; // namespace milvus::futures \ No newline at end of file +}; // namespace milvus::futures diff --git a/internal/core/src/futures/Executor.h b/internal/core/src/futures/Executor.h index 6f6d4d1a4c..31e498dba0 100644 --- a/internal/core/src/futures/Executor.h +++ b/internal/core/src/futures/Executor.h @@ -33,4 +33,7 @@ getSearchCPUExecutor(); folly::CPUThreadPoolExecutor* getLoadCPUExecutor(); +folly::CPUThreadPoolExecutor* +getJsonStatsBuildExecutor(); + }; // namespace milvus::futures diff --git a/internal/core/src/futures/FutureTest.cpp b/internal/core/src/futures/FutureTest.cpp index 82e60dbefc..e8768aaa1b 100644 --- a/internal/core/src/futures/FutureTest.cpp +++ b/internal/core/src/futures/FutureTest.cpp @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -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 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(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(original_capacity)); + EXPECT_EQ(status.error_code, milvus::Success); + EXPECT_EQ(budget.CapacityBytes(), original_capacity); +} + TEST(Futures, LeakyResult) { { LeakyResult leaky_result; @@ -225,4 +277,4 @@ TEST(Futures, Future) { ASSERT_EQ(s.error_code, milvus::FollyCancel); free((char*)(s.error_msg)); } -} \ No newline at end of file +} diff --git a/internal/core/src/futures/future_c.cpp b/internal/core/src/futures/future_c.cpp index 3614fb5ba6..e7eb2bf247 100644 --- a/internal/core/src/futures/future_c.cpp +++ b/internal/core/src/futures/future_c.cpp @@ -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 +#include + +#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(std::max(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 +} diff --git a/internal/core/src/futures/future_c.h b/internal/core/src/futures/future_c.h index d9f8636d71..498bce5e11 100644 --- a/internal/core/src/futures/future_c.h +++ b/internal/core/src/futures/future_c.h @@ -14,6 +14,7 @@ #include "common/common_type_c.h" #include "future_c_types.h" #include +#include #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 diff --git a/internal/core/src/index/json_stats/JsonKeyStats.cpp b/internal/core/src/index/json_stats/JsonKeyStats.cpp index 8be54189e3..7983c574d0 100644 --- a/internal/core/src/index/json_stats/JsonKeyStats.cpp +++ b/internal/core/src/index/json_stats/JsonKeyStats.cpp @@ -14,17 +14,21 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include -#include +#include +#include #include #include #include #include #include #include +#include +#include +#include #include +#include #include -#include "segcore/default_fs.h" +#include #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 infos; +}; + +template +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(bit_width)))); + } + return buffer_bytes; +} struct JsonStatsParquetMetadata { std::shared_ptr schema; @@ -205,6 +297,93 @@ GetJsonStatsFieldsFromSchema(const std::shared_ptr& schema) { } // namespace +std::vector +CreateJsonStatsRowRanges(const std::vector& 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 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(raw_value); + input_bytes = SaturatingAdd(input_bytes, json->data().size()); + }); + + const auto row_count = static_cast(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(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(ctx); disk_file_manager_ = std::make_shared(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& field_datas, return infos; } +std::map +JsonKeyStats::CollectKeyInfo(const std::vector& row_ranges, + bool nullable) { + std::map 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(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> 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 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( + 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 JsonKeyStats::ClassifyJsonKeyLayoutType( const std::map& infos) { @@ -551,7 +828,7 @@ void JsonKeyStats::AddKeyStats(const std::vector& path, JSONType type, const std::string& value, - std::map& values) { + std::map& 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& path, - std::map& values) { + std::map& 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& 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 paths; - std::map 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 values; + if (!ParseJsonForBuildStats(json_str, values)) { + BuildKeyStatsForNullRow(); + return; + } + DomNode root; std::set 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& field_datas, +JsonKeyStats::BuildKeyStats(const std::vector& 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(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( + 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(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& 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(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> 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(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(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 values; + if (!ParseJsonForBuildStats(json_str, values)) { + append_null_row(); + return; + } + + DomNode root; + std::set 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( + 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(global_row), offset)); + } + + auto shared_builder = + std::static_pointer_cast(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> 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 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)); + } + for (const auto& [key, postings] : bson_postings) { + materialized_bytes = SaturatingAdd( + materialized_bytes, + sizeof(std::pair>) + + 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& 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(1, executor->numThreads())); + auto& memory_budget = JsonStatsBuildMemoryBudget::GetInstance(); + + struct ActiveMaterializeTask { + std::shared_ptr> accounted_bytes; + folly::Future future; + }; + + std::vector 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> accounted_bytes; + try { + accounted_bytes = + std::make_shared>(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& 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(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& 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), diff --git a/internal/core/src/index/json_stats/JsonKeyStats.h b/internal/core/src/index/json_stats/JsonKeyStats.h index b24691f7ea..e5a4a38453 100644 --- a/internal/core/src/index/json_stats/JsonKeyStats.h +++ b/internal/core/src/index/json_stats/JsonKeyStats.h @@ -29,6 +29,7 @@ #include #include #include +#include #include #include #include @@ -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 slices; +}; + +std::vector +CreateJsonStatsRowRanges(const std::vector& field_datas, + int64_t max_rows_per_range); + class JsonKeyStats : public ScalarIndex { 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 { return ss.str(); } + // Keep the original serial implementation for differential tests. std::map CollectKeyInfo(const std::vector& field_datas, bool nullable); + std::map + CollectKeyInfo(const std::vector& row_ranges, + bool nullable); + void TraverseJsonForStats(const char* json, jsmntok* tokens, @@ -496,8 +523,30 @@ class JsonKeyStats : public ScalarIndex { std::map ClassifyJsonKeyLayoutType(const std::map& infos); + // Keep the original serial implementation for differential tests. void - BuildKeyStats(const std::vector& field_datas, bool nullable); + BuildKeyStats(const std::vector& row_ranges, + bool nullable); + + void + BuildKeyStatsParallel(const std::vector& row_ranges, + bool nullable); + + static size_t + EstimateJsonStatsMaterializeReservationBytes(const JsonStatsRowRange& range, + const arrow::Schema& schema); + + struct MaterializedChunk { + std::shared_ptr record_batch; + std::map> bson_postings; + size_t memory_bytes; + }; + + MaterializedChunk + MaterializeKeyStatsRange( + const JsonStatsRowRange& range, + bool nullable, + const std::shared_ptr& schema) const; void BuildKeyStatsForRow(std::string_view json_str, uint32_t row_id); @@ -524,22 +573,26 @@ class JsonKeyStats : public ScalarIndex { AddKeyStats(const std::vector& path, JSONType type, const std::string& value, - std::map& values); + std::map& values) const; + + bool + ParseJsonForBuildStats(std::string_view json_str, + std::map& values) const; void TraverseJsonForBuildStats(const char* json, jsmntok* tokens, int& index, std::vector& path, - std::map& values); + std::map& 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 { } 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 { // 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; }; diff --git a/internal/core/src/index/json_stats/JsonStatsBuildMemoryBudget.h b/internal/core/src/index/json_stats/JsonStatsBuildMemoryBudget.h new file mode 100644 index 0000000000..f3c5fda1dc --- /dev/null +++ b/internal/core/src/index/json_stats/JsonStatsBuildMemoryBudget.h @@ -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 +#include +#include +#include + +#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 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 lock(mu_); + if (!CanAcquireLocked(bytes)) { + return false; + } + inflight_bytes_ += bytes; + return true; + } + + void + Release(size_t bytes) { + { + std::lock_guard 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 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::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 lock(mu_); + return capacity_bytes_; + } + + size_t + InflightBytes() const { + std::lock_guard lock(mu_); + return inflight_bytes_; + } + + void + SetCapacityBytes(size_t bytes) { + { + std::lock_guard 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 diff --git a/internal/core/src/index/json_stats/bson_inverted.cpp b/internal/core/src/index/json_stats/bson_inverted.cpp index 35f86c2bff..25ea15e10c 100644 --- a/internal/core/src/index/json_stats/bson_inverted.cpp +++ b/internal/core/src/index/json_stats/bson_inverted.cpp @@ -76,6 +76,25 @@ BsonInvertedIndex::AddRecord(const std::string& key, EncodeInvertedIndexValue(row_id, offset)); } +void +BsonInvertedIndex::AddRecords(const std::string& key, + std::vector 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) { diff --git a/internal/core/src/index/json_stats/bson_inverted.h b/internal/core/src/index/json_stats/bson_inverted.h index b7a92f8335..2bdaa6e185 100644 --- a/internal/core/src/index/json_stats/bson_inverted.h +++ b/internal/core/src/index/json_stats/bson_inverted.h @@ -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 records); + void BuildIndex(); diff --git a/internal/core/src/index/json_stats/parquet_writer.cpp b/internal/core/src/index/json_stats/parquet_writer.cpp index 99ff85fed9..d2b08c82cb 100644 --- a/internal/core/src/index/json_stats/parquet_writer.cpp +++ b/internal/core/src/index/json_stats/parquet_writer.cpp @@ -18,7 +18,9 @@ #include #include +#include #include +#include #include #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& 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(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(batch->num_rows()); + return arrow::Status::OK(); +} + +arrow::Status +AppendJsonStatsValueToBuilder( const std::string& value, const std::shared_ptr& 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: {}", diff --git a/internal/core/src/index/json_stats/parquet_writer.h b/internal/core/src/index/json_stats/parquet_writer.h index 81a7a36195..4b9c4a30b7 100644 --- a/internal/core/src/index/json_stats/parquet_writer.h +++ b/internal/core/src/index/json_stats/parquet_writer.h @@ -101,6 +101,11 @@ struct TableStatsInfo { std::map column_map; }; +arrow::Status +AppendJsonStatsValueToBuilder( + const std::string& value, + const std::shared_ptr& 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& batch); + arrow::Status Flush(); @@ -220,14 +228,15 @@ class JsonStatsParquetWriter { return total_size_; } + const std::shared_ptr& + GetSchema() const { + return schema_; + } + void UpdatePathSizeMap(const std::vector>& arrays); private: - arrow::Status - AppendDataToBuilder(const std::string& value, - const std::shared_ptr& builder); - // init info std::shared_ptr schema_{nullptr}; diff --git a/internal/core/src/index/json_stats/utils.cpp b/internal/core/src/index/json_stats/utils.cpp index 927b64ea18..717d26ee05 100644 --- a/internal/core/src/index/json_stats/utils.cpp +++ b/internal/core/src/index/json_stats/utils.cpp @@ -151,7 +151,8 @@ CreateArrowBuilders(const std::map& 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& column_map) { return {std::move(builders), std::move(builders_map)}; } +std::pair>, + std::map>> +CreateArrowBuildersForColumns(const std::set& column_keys) { + std::vector> builders; + builders.reserve(column_keys.size() + 1); + std::map> 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 CreateArrowSchema(const std::map& column_map) { std::vector> fields; @@ -315,4 +331,4 @@ JsonStatsMeta::DeserializeToKeyFieldMap(const std::string& json_str) { return key_field_map; } -} // namespace milvus::index \ No newline at end of file +} // namespace milvus::index diff --git a/internal/core/src/index/json_stats/utils.h b/internal/core/src/index/json_stats/utils.h index 2a6dc3aa7b..a58e012652 100644 --- a/internal/core/src/index/json_stats/utils.h +++ b/internal/core/src/index/json_stats/utils.h @@ -19,6 +19,7 @@ #include #include "common/FastMem.h" #include +#include #include #include #include @@ -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::map>> CreateArrowBuilders(const std::map& 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::map>> +CreateArrowBuildersForColumns(const std::set& column_keys); + std::shared_ptr CreateArrowSchema(const std::map& column_map); diff --git a/internal/core/unittest/test_json_stats/test_json_key_stats.cpp b/internal/core/unittest/test_json_stats/test_json_key_stats.cpp index a3bb98a2e9..faf0b2d2dc 100644 --- a/internal/core/unittest/test_json_stats/test_json_key_stats.cpp +++ b/internal/core/unittest/test_json_stats/test_json_key_stats.cpp @@ -18,11 +18,15 @@ #include #include #include +#include #include +#include +#include #include #include #include #include +#include #include #include #include @@ -45,14 +49,18 @@ #include "common/Types.h" #include "common/bson_view.h" #include "common/protobuf_utils.h" +#include "futures/Executor.h" +#include "futures/future_c.h" #include "gtest/gtest.h" #include "index/IndexInfo.h" #include "index/IndexStats.h" #include "index/Meta.h" #include "index/json_stats/JsonKeyStats.h" +#include "index/json_stats/JsonStatsBuildMemoryBudget.h" #include "index/json_stats/bson_inverted.h" #include "index/json_stats/utils.h" #include "indexbuilder/IndexCreatorBase.h" +#include "milvus-storage/format/parquet/file_reader.h" #include "pb/common.pb.h" #include "pb/schema.pb.h" #include "segcore/storagev2translator/JsonStatsTranslator.h" @@ -73,6 +81,79 @@ using namespace milvus::indexbuilder; using namespace milvus; using namespace milvus::index; +class BuildKeyStatsAccessor { + public: + static void + Prepare(JsonKeyStats& stats, + const std::map& key_types, + int64_t num_rows) { + stats.num_rows_ = num_rows; + stats.key_types_ = key_types; + stats.shared_keys_.clear(); + stats.column_keys_.clear(); + for (const auto& [key, layout] : key_types) { + if (layout == JsonKeyLayoutType::SHARED) { + stats.shared_keys_.insert(key); + } else { + stats.column_keys_.insert(key); + } + } + + auto remote_prefix = + stats.disk_file_manager_->GetRemoteJsonStatsShreddingPrefix(); + auto writer_context = + ParquetWriterFactory::CreateContext(key_types, remote_prefix); + stats.parquet_writer_->Init(std::move(writer_context)); + } + + static void + CallLegacy(JsonKeyStats& stats, + const std::vector& row_ranges, + bool nullable) { + stats.BuildKeyStats(row_ranges, nullable); + } + + static void + CallParallel(JsonKeyStats& stats, + const std::vector& row_ranges, + bool nullable) { + stats.BuildKeyStatsParallel(row_ranges, nullable); + } + + static size_t + EstimateMaterializeReservationBytes(const JsonStatsRowRange& range, + const arrow::Schema& schema) { + return JsonKeyStats::EstimateJsonStatsMaterializeReservationBytes( + range, schema); + } + + static void + Finish(JsonKeyStats& stats) { + auto status = stats.parquet_writer_->Close(); + AssertInfo(status.ok(), + "failed to close json stats test writer: {}", + status.ToString()); + stats.bson_inverted_index_->BuildIndex(); + } + + static std::vector + GetParquetFiles(const JsonKeyStats& stats) { + std::vector files; + for (const auto& [path, size] : + stats.parquet_writer_->GetPathsToSize()) { + if (size > 0) { + files.push_back(path); + } + } + return files; + } + + static IndexStatsPtr + UploadBsonIndex(JsonKeyStats& stats) { + return stats.bson_inverted_index_->UploadIndex(); + } +}; + int64_t GenerateRandomInt64(int64_t min, int64_t max) { static std::random_device rd; @@ -176,6 +257,1099 @@ GenerateJsons(int size) { return jsons; } +namespace { + +milvus::FieldDataPtr +CreateMaterializeJsonFieldData(const std::vector& json_rows) { + auto field_data = milvus::storage::CreateFieldData( + milvus::DataType::JSON, milvus::DataType::NONE, false); + if (json_rows.empty()) { + return field_data; + } + + std::vector 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 +CreateNullableMaterializeJsonFieldData( + const std::vector& json_rows, + const std::vector& valid_rows) { + AssertInfo(json_rows.size() == valid_rows.size(), + "json rows and validity rows must have the same size"); + auto field_data = milvus::storage::CreateFieldData( + milvus::DataType::JSON, milvus::DataType::NONE, true); + if (json_rows.empty()) { + return field_data; + } + + std::vector rows; + rows.reserve(json_rows.size()); + for (const auto& json : json_rows) { + rows.emplace_back(simdjson::padded_string(json)); + } + + std::vector 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(1U << (i % 8)); + } + } + field_data->FillFieldData(rows.data(), valid_bitmap.data(), rows.size(), 0); + return field_data; +} + +struct MaterializeTestIndex { + std::shared_ptr chunk_manager; + milvus_storage::ArrowFileSystemPtr fs; + storage::FileManagerContext context; + std::unique_ptr stats; +}; + +MaterializeTestIndex +CreateMaterializeTestIndex(bool nullable, int64_t write_batch_size) { + static std::atomic next_id{900000}; + const auto id = next_id.fetch_add(1); + + proto::schema::FieldSchema field_schema; + field_schema.set_data_type(proto::schema::DataType::JSON); + field_schema.set_nullable(nullable); + storage::FieldDataMeta field_meta{11, 12, id, 101, field_schema}; + storage::IndexMeta index_meta{id, 101, id, 1}; + + storage::StorageConfig storage_config; + storage_config.storage_type = "local"; + storage_config.root_path = TestLocalPath; + auto chunk_manager = storage::CreateChunkManager(storage_config); + auto fs = storage::InitArrowFileSystem(storage_config); + storage::FileManagerContext ctx(field_meta, index_meta, chunk_manager, fs); + auto stats = + std::make_unique(ctx, false, 1024, 0.3, write_batch_size); + + return MaterializeTestIndex{ + chunk_manager, + fs, + ctx, + std::move(stats), + }; +} + +struct MaterializedTable { + std::shared_ptr table; + std::vector row_group_rows; +}; + +MaterializedTable +ReadMaterializedTable(const milvus_storage::ArrowFileSystemPtr& fs, + const std::string& file) { + auto reader_result = milvus_storage::FileRowGroupReader::Make(fs, file); + AssertInfo(reader_result.ok(), + "failed to create json stats parquet test reader: {}", + reader_result.status().ToString()); + auto reader = reader_result.ValueOrDie(); + auto row_groups = reader->file_metadata()->GetRowGroupMetadataVector(); + AssertInfo(row_groups.size() > 0, + "json stats parquet test file has no row groups"); + + auto status = reader->SetRowGroupOffsetAndCount( + 0, static_cast(row_groups.size())); + AssertInfo(status.ok(), + "failed to select json stats parquet row groups: {}", + status.ToString()); + + std::vector> tables; + std::vector row_group_rows; + tables.reserve(row_groups.size()); + row_group_rows.reserve(row_groups.size()); + for (size_t i = 0; i < row_groups.size(); ++i) { + std::shared_ptr table; + status = reader->ReadNextRowGroup(&table); + AssertInfo(status.ok(), + "failed to read json stats parquet row group: {}", + status.ToString()); + AssertInfo(table != nullptr, + "json stats parquet reader returned a null row group"); + tables.push_back(std::move(table)); + row_group_rows.push_back(row_groups.Get(i).row_num()); + } + status = reader->Close(); + AssertInfo(status.ok(), + "failed to close json stats parquet test reader: {}", + status.ToString()); + + auto concatenate_result = arrow::ConcatenateTables(tables); + AssertInfo(concatenate_result.ok(), + "failed to concatenate json stats parquet tables: {}", + concatenate_result.status().ToString()); + auto combine_result = concatenate_result.ValueOrDie()->CombineChunks(); + AssertInfo(combine_result.ok(), + "failed to combine json stats parquet chunks: {}", + combine_result.status().ToString()); + return MaterializedTable{ + combine_result.ValueOrDie(), + std::move(row_group_rows), + }; +} + +struct MaterializeSnapshot { + MaterializedTable parquet; + std::map>> postings; +}; + +MaterializeSnapshot +BuildMaterializeSnapshot(const std::vector& row_ranges, + const std::map& key_types, + const std::vector& shared_paths, + bool nullable, + bool parallel, + int64_t write_batch_size = 5) { + int64_t num_rows = 0; + for (const auto& range : row_ranges) { + num_rows += range.row_count; + } + + auto test_index = CreateMaterializeTestIndex(nullable, write_batch_size); + BuildKeyStatsAccessor::Prepare(*test_index.stats, key_types, num_rows); + if (parallel) { + BuildKeyStatsAccessor::CallParallel( + *test_index.stats, row_ranges, nullable); + } else { + BuildKeyStatsAccessor::CallLegacy( + *test_index.stats, row_ranges, nullable); + } + BuildKeyStatsAccessor::Finish(*test_index.stats); + + auto files = BuildKeyStatsAccessor::GetParquetFiles(*test_index.stats); + if (num_rows == 0) { + AssertInfo(files.empty(), + "empty json stats build unexpectedly produced {} parquet " + "files", + files.size()); + return MaterializeSnapshot{ + MaterializedTable{nullptr, {}}, + {}, + }; + } + AssertInfo(files.size() == 1, + "expected one json stats parquet file, got {}", + files.size()); + + MaterializeSnapshot snapshot{ + ReadMaterializedTable(test_index.fs, files.front()), + {}, + }; + auto bson_stats = BuildKeyStatsAccessor::UploadBsonIndex(*test_index.stats); + auto load_disk_file_manager = + std::make_shared(test_index.context); + BsonInvertedIndex loaded_bson_index(load_disk_file_manager); + loaded_bson_index.LoadIndex(bson_stats->GetIndexFiles(), + milvus::proto::common::LoadPriority::HIGH, + false); + for (const auto& path : shared_paths) { + loaded_bson_index.TermQuery( + path, + [&](const uint32_t* row_ids, + const uint32_t* offsets, + int64_t count) { + for (int64_t i = 0; i < count; ++i) { + snapshot.postings[path].emplace_back(row_ids[i], + offsets[i]); + } + }); + } + return snapshot; +} + +void +ExpectMaterializeSnapshotsEqual(const MaterializeSnapshot& expected, + const MaterializeSnapshot& actual, + bool compare_row_groups = false) { + ASSERT_EQ(expected.parquet.table == nullptr, + actual.parquet.table == nullptr); + if (expected.parquet.table == nullptr) { + EXPECT_EQ(actual.postings, expected.postings); + return; + } + + EXPECT_TRUE(expected.parquet.table->Equals(*actual.parquet.table, false)) + << "expected:\n" + << expected.parquet.table->ToString() << "\nactual:\n" + << actual.parquet.table->ToString(); + ASSERT_EQ(expected.parquet.table->num_columns(), + actual.parquet.table->num_columns()); + for (int column = 0; column < expected.parquet.table->num_columns(); + ++column) { + EXPECT_TRUE(expected.parquet.table->field(column)->Equals( + *actual.parquet.table->field(column), true)); + } + if (compare_row_groups) { + EXPECT_EQ(actual.parquet.row_group_rows, + expected.parquet.row_group_rows); + } + EXPECT_EQ(actual.postings, expected.postings); +} + +struct MixedMaterializeDataset { + std::vector field_datas; + std::map key_types; + std::vector shared_paths; + int64_t row_count; + int64_t expected_dense_postings; +}; + +MixedMaterializeDataset +CreateMixedMaterializeDataset() { + constexpr int64_t kRowCount = 97; + std::vector rows; + rows.reserve(kRowCount); + for (int64_t i = 0; i < kRowCount; ++i) { + nlohmann::json row; + row["id"] = i % 2 == 0 ? i : -i; + row["flag"] = i % 3 == 0; + row["mixed"] = i % 2 == 0 ? nlohmann::json(i * 10) + : nlohmann::json(fmt::format("mixed-{}", i)); + row["dense"] = fmt::format("dense-{}-雪", i); + row["nested"]["value"] = i * 100; + row["a/b"]["~key"] = fmt::format("escaped-{}", i); + if (i % 2 == 0) { + row["text"] = i % 10 == 0 ? "" : fmt::format("text-{}\nline", i); + } + if (i % 3 != 0) { + row["score"] = static_cast(i) + 0.125; + } + if (i % 4 == 0) { + row["arr"] = nlohmann::json::array( + {i, fmt::format("item-{}", i), i % 8 == 0}); + } + if (i % 10 == 0) { + row["empty_obj"] = nlohmann::json::object(); + } + if (i % 13 == 0) { + row["nullv"] = nullptr; + } + rows.push_back(row.dump()); + } + + for (auto row : {5, 44, 96}) { + rows[row].clear(); + } + + std::vector valid_rows(kRowCount, true); + constexpr int64_t kNullableBegin = 19; + constexpr int64_t kNullableEnd = 67; + for (int64_t row = kNullableBegin; row < kNullableEnd; ++row) { + if (row % 11 == 0) { + valid_rows[row] = false; + } + } + + int64_t expected_dense_postings = 0; + for (int64_t row = 0; row < kRowCount; ++row) { + if (valid_rows[row] && !rows[row].empty()) { + ++expected_dense_postings; + } + } + + std::vector first(rows.begin(), rows.begin() + kNullableBegin); + std::vector nullable(rows.begin() + kNullableBegin, + rows.begin() + kNullableEnd); + std::vector nullable_valid(valid_rows.begin() + kNullableBegin, + valid_rows.begin() + kNullableEnd); + std::vector last(rows.begin() + kNullableEnd, rows.end()); + + return MixedMaterializeDataset{ + { + CreateMaterializeJsonFieldData(first), + CreateMaterializeJsonFieldData({}), + CreateNullableMaterializeJsonFieldData(nullable, nullable_valid), + CreateMaterializeJsonFieldData(last), + }, + { + {JsonKey("/id", JSONType::INT64), JsonKeyLayoutType::TYPED}, + {JsonKey("/flag", JSONType::BOOL), + JsonKeyLayoutType::TYPED_NOT_ALL}, + {JsonKey("/text", JSONType::STRING), + JsonKeyLayoutType::TYPED_NOT_ALL}, + {JsonKey("/score", JSONType::DOUBLE), + JsonKeyLayoutType::TYPED_NOT_ALL}, + {JsonKey("/arr", JSONType::ARRAY), + JsonKeyLayoutType::TYPED_NOT_ALL}, + {JsonKey("/mixed", JSONType::INT64), JsonKeyLayoutType::DYNAMIC}, + {JsonKey("/mixed", JSONType::STRING), JsonKeyLayoutType::DYNAMIC}, + {JsonKey("/dense", JSONType::STRING), JsonKeyLayoutType::SHARED}, + {JsonKey("/nested/value", JSONType::INT64), + JsonKeyLayoutType::SHARED}, + {JsonKey("/a~1b/~0key", JSONType::STRING), + JsonKeyLayoutType::SHARED}, + {JsonKey("/empty_obj", JSONType::OBJECT), + JsonKeyLayoutType::SHARED}, + {JsonKey("/nullv", JSONType::NONE), JsonKeyLayoutType::SHARED}, + }, + { + "/dense", + "/nested", + "/nested/value", + "/a~1b", + "/a~1b/~0key", + "/empty_obj", + "/nullv", + }, + kRowCount, + expected_dense_postings, + }; +} + +} // namespace + +TEST(JsonStatsMemoryBudgetTest, CapacityUpdateWakesBlockedBuild) { + auto& budget = JsonStatsBuildMemoryBudget::GetInstance(); + ASSERT_EQ(budget.InflightBytes(), 0); + auto original_capacity = budget.CapacityBytes(); + auto restore_capacity = folly::makeGuard([&]() { + executor_set_json_stats_build_max_inflight_bytes( + static_cast(original_capacity)); + }); + + executor_set_json_stats_build_max_inflight_bytes(1); + EXPECT_EQ(budget.CapacityBytes(), 1); + budget.Acquire(1); + auto release_budget = folly::makeGuard([&]() { budget.Release(1); }); + + auto waiter = std::async(std::launch::async, [&]() { + budget.Acquire(1); + budget.Release(1); + }); + EXPECT_EQ(waiter.wait_for(std::chrono::milliseconds(50)), + std::future_status::timeout); + + executor_set_json_stats_build_max_inflight_bytes(2); + EXPECT_EQ(budget.CapacityBytes(), 2); + EXPECT_EQ(waiter.wait_for(std::chrono::seconds(5)), + std::future_status::ready); + waiter.get(); + + executor_set_json_stats_build_max_inflight_bytes(-1); + EXPECT_EQ(budget.CapacityBytes(), 0); +} + +TEST(JsonStatsMemoryBudgetTest, + MaterializeReservationAccountsForSparseArrowColumns) { + constexpr int64_t kRowCount = 16 * 1024; + constexpr size_t kShreddingColumnCount = 1024; + constexpr size_t kBudgetBytes = 512UL * 1024 * 1024; + + auto field_data = CreateMaterializeJsonFieldData( + std::vector(kRowCount, "{}")); + auto ranges = CreateJsonStatsRowRanges({field_data}, kRowCount); + ASSERT_EQ(ranges.size(), 1); + + std::vector> fields; + fields.reserve(kShreddingColumnCount + 1); + for (size_t i = 0; i < kShreddingColumnCount; ++i) { + fields.push_back( + arrow::field(fmt::format("int64_{}", i), arrow::int64(), true)); + } + fields.push_back(arrow::field("shared", arrow::binary(), true)); + auto schema = arrow::schema(std::move(fields)); + + const auto reservation = + BuildKeyStatsAccessor::EstimateMaterializeReservationBytes( + ranges.front(), *schema); + const auto validity_bytes = static_cast((kRowCount + 7) / 8); + const auto shredding_buffer_bytes = + kShreddingColumnCount * + (static_cast(kRowCount) * sizeof(int64_t) + validity_bytes); + const auto shared_buffer_bytes = + validity_bytes + (static_cast(kRowCount) + 1) * + sizeof(arrow::BinaryType::offset_type); + + EXPECT_GE(reservation, shredding_buffer_bytes + shared_buffer_bytes); + EXPECT_GT(reservation, kBudgetBytes / 4); +} + +TEST(JsonStatsMaterializeTest, NonPositiveWriterBatchUsesDefault) { + constexpr int64_t kRowCount = 16 * 1024 + 1; + constexpr int64_t kDefaultWriteBatchRows = 81920; + + auto original_thread_num = + milvus::futures::getJsonStatsBuildExecutor()->numThreads(); + auto restore_threads = folly::makeGuard([&]() { + executor_set_json_stats_build_thread_num( + static_cast(original_thread_num)); + }); + + const std::string value(512, 'x'); + std::vector rows; + rows.reserve(kRowCount); + for (int64_t row = 0; row < kRowCount; ++row) { + rows.push_back(fmt::format(R"({{"int":{},"text":"{}"}})", row, value)); + } + auto field_data = CreateMaterializeJsonFieldData(rows); + + auto build = [&](int64_t write_batch_size) { + // Pass the request value directly to the native constructor, as an + // older coordinator can bypass the Go configuration formatter. + auto test_index = CreateMaterializeTestIndex(false, write_batch_size); + test_index.stats->BuildWithFieldData({field_data}, false); + auto files = BuildKeyStatsAccessor::GetParquetFiles(*test_index.stats); + EXPECT_EQ(files.size(), 1); + return ReadMaterializedTable(test_index.fs, files.at(0)); + }; + + for (int worker_count : {1, 4}) { + SCOPED_TRACE(::testing::Message() << "worker_count=" << worker_count); + executor_set_json_stats_build_thread_num(worker_count); + auto expected = build(kDefaultWriteBatchRows); + ASSERT_NE(expected.table, nullptr); + ASSERT_EQ(expected.table->num_rows(), kRowCount); + + for (int64_t batch_size : + {int64_t{0}, int64_t{-1}, std::numeric_limits::min()}) { + SCOPED_TRACE(::testing::Message() << "batch_size=" << batch_size); + auto actual = build(batch_size); + ASSERT_NE(actual.table, nullptr); + EXPECT_TRUE(expected.table->Equals(*actual.table)); + EXPECT_EQ(expected.row_group_rows, actual.row_group_rows); + } + } +} + +TEST(JsonStatsMaterializeTest, SmallWriterBatchPreservesParquetRowGroupSizing) { + constexpr int64_t kMaterializeRangeRows = 16 * 1024; + constexpr int64_t kWriteBatchRows = 1024; + constexpr int kWorkerCount = 4; + + auto original_thread_num = + milvus::futures::getJsonStatsBuildExecutor()->numThreads(); + auto restore_threads = folly::makeGuard([&]() { + executor_set_json_stats_build_thread_num( + static_cast(original_thread_num)); + }); + executor_set_json_stats_build_thread_num(kWorkerCount); + + // One row beyond the fixed 16,384-row range forces the production build + // down the parallel materialization path. A 512-byte value makes correctly + // accounted Parquet row groups span multiple 1,024-row writer batches. + const std::string value(512, 'x'); + const std::string json = fmt::format(R"({{"text":"{}"}})", value); + auto field_data = CreateMaterializeJsonFieldData( + std::vector(kMaterializeRangeRows + 1, json)); + auto test_index = CreateMaterializeTestIndex(false, kWriteBatchRows); + + test_index.stats->BuildWithFieldData({field_data}, false); + + auto files = BuildKeyStatsAccessor::GetParquetFiles(*test_index.stats); + ASSERT_EQ(files.size(), 1); + auto materialized = ReadMaterializedTable(test_index.fs, files.front()); + ASSERT_NE(materialized.table, nullptr); + EXPECT_EQ(materialized.table->num_rows(), kMaterializeRangeRows + 1); + ASSERT_FALSE(materialized.row_group_rows.empty()); + EXPECT_GT(*std::max_element(materialized.row_group_rows.begin(), + materialized.row_group_rows.end()), + kWriteBatchRows); +} + +TEST(JsonStatsMaterializeTest, + ParallelMatchesLegacyAcrossRangesAndWriterBatches) { + std::map key_types = { + {JsonKey("/int", JSONType::INT64), JsonKeyLayoutType::TYPED_NOT_ALL}, + {JsonKey("/double", JSONType::DOUBLE), + JsonKeyLayoutType::TYPED_NOT_ALL}, + {JsonKey("/bool", JSONType::BOOL), JsonKeyLayoutType::TYPED_NOT_ALL}, + {JsonKey("/text", JSONType::STRING), JsonKeyLayoutType::TYPED_NOT_ALL}, + {JsonKey("/arr", JSONType::ARRAY), JsonKeyLayoutType::TYPED_NOT_ALL}, + {JsonKey("/mixed", JSONType::INT64), JsonKeyLayoutType::DYNAMIC}, + {JsonKey("/mixed", JSONType::STRING), JsonKeyLayoutType::DYNAMIC}, + {JsonKey("/rare", JSONType::STRING), JsonKeyLayoutType::SHARED}, + {JsonKey("/obj", JSONType::OBJECT), JsonKeyLayoutType::SHARED}, + {JsonKey("/nullv", JSONType::NONE), JsonKeyLayoutType::SHARED}, + {JsonKey("/nested/value", JSONType::DOUBLE), JsonKeyLayoutType::SHARED}, + }; + + std::vector field_datas = { + CreateMaterializeJsonFieldData({ + R"({"int":1,"double":1.5,"bool":true,"text":"one","arr":[1,2],"mixed":1,"rare":"first"})", + "", + }), + CreateNullableMaterializeJsonFieldData( + { + R"({"int":2,"double":2.5,"bool":false,"text":"two","arr":[],"mixed":"two","nested":{"value":2.25}})", + R"({"int":300,"double":300.5,"bool":true,"text":"invalid-row","arr":[300],"mixed":300,"rare":"ignored"})", + R"({})", + R"({"int":5,"double":5.5,"bool":false,"text":"escaped\nvalue","arr":[{"x":1}],"mixed":5,"obj":{},"nullv":null})", + R"({"int":6,"double":6.5,"bool":true,"text":"six","arr":[6],"mixed":"six","rare":"middle"})", + }, + {true, false, true, true, true}), + CreateMaterializeJsonFieldData({ + R"({"int":7,"double":7.5,"bool":false,"text":"seven","arr":[7,8],"mixed":7})", + R"({"int":8,"double":8.5,"bool":true,"text":"eight","arr":[],"mixed":"eight","rare":"last","nested":{"value":8.25}})", + R"({"int":9,"double":9.5,"bool":false,"text":"nine","arr":[9],"mixed":9,"obj":{}})", + R"({"int":10,"double":10.5,"bool":true,"text":"ten","arr":[10],"mixed":"ten","nullv":null})", + }), + }; + auto row_ranges = CreateJsonStatsRowRanges(field_datas, 3); + ASSERT_EQ(row_ranges.size(), 4); + ASSERT_EQ(row_ranges[0].slices.size(), 2); + ASSERT_EQ(row_ranges[2].slices.size(), 2); + + const std::vector shared_paths = { + "/rare", "/obj", "/nullv", "/nested", "/nested/value"}; + auto original_thread_num = + milvus::futures::getJsonStatsBuildExecutor()->numThreads(); + auto restore_threads = folly::makeGuard([&]() { + executor_set_json_stats_build_thread_num( + static_cast(original_thread_num)); + }); + + auto legacy = BuildMaterializeSnapshot( + row_ranges, key_types, shared_paths, true, false); + EXPECT_EQ(legacy.parquet.table->num_rows(), 11); + EXPECT_FALSE(legacy.parquet.row_group_rows.empty()); + + for (int worker_count : {1, 2, 4}) { + SCOPED_TRACE(::testing::Message() << "worker_count=" << worker_count); + executor_set_json_stats_build_thread_num(worker_count); + auto parallel = BuildMaterializeSnapshot( + row_ranges, key_types, shared_paths, true, true); + ExpectMaterializeSnapshotsEqual(legacy, parallel, true); + } +} + +TEST(JsonStatsMaterializeTest, + ParallelPreservesNonNulTerminatedJsonViewLength) { + const std::string json = R"({"int":7,"shared":"value"})"; + std::string buffer = json; + buffer.append(simdjson::SIMDJSON_PADDING, '\xAB'); + + auto field_data = std::dynamic_pointer_cast( + 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())}); + + std::map key_types = { + {JsonKey("/int", JSONType::INT64), JsonKeyLayoutType::TYPED}, + {JsonKey("/shared", JSONType::STRING), JsonKeyLayoutType::SHARED}, + }; + auto ranges = CreateJsonStatsRowRanges({field_data}, 1); + auto legacy = BuildMaterializeSnapshot( + ranges, key_types, {"/shared"}, false, false, 1); + auto parallel = BuildMaterializeSnapshot( + ranges, key_types, {"/shared"}, false, true, 1); + + ASSERT_NE(parallel.parquet.table, nullptr); + EXPECT_EQ(parallel.parquet.table->num_rows(), 1); + ExpectMaterializeSnapshotsEqual(legacy, parallel, true); +} + +TEST(JsonStatsMaterializeTest, ParallelMatchesLegacyForEmptyInput) { + std::vector field_datas = { + CreateMaterializeJsonFieldData({}), + CreateNullableMaterializeJsonFieldData({}, {}), + }; + auto row_ranges = CreateJsonStatsRowRanges(field_datas, 3); + ASSERT_TRUE(row_ranges.empty()); + + std::map key_types = { + {JsonKey("/int", JSONType::INT64), JsonKeyLayoutType::TYPED}, + {JsonKey("/rare", JSONType::STRING), JsonKeyLayoutType::SHARED}, + }; + auto legacy = + BuildMaterializeSnapshot(row_ranges, key_types, {}, true, false, 3); + ASSERT_EQ(legacy.parquet.table, nullptr); + + auto original_thread_num = + milvus::futures::getJsonStatsBuildExecutor()->numThreads(); + auto restore_threads = folly::makeGuard([&]() { + executor_set_json_stats_build_thread_num( + static_cast(original_thread_num)); + }); + 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 = + BuildMaterializeSnapshot(row_ranges, key_types, {}, true, true, 1); + ExpectMaterializeSnapshotsEqual(legacy, parallel); + } +} + +TEST(JsonStatsMaterializeTest, + ParallelMatchesLegacyForNullLikeRowsAndEmptyFieldData) { + std::map key_types = { + {JsonKey("/int", JSONType::INT64), JsonKeyLayoutType::TYPED}, + {JsonKey("/text", JSONType::STRING), JsonKeyLayoutType::TYPED_NOT_ALL}, + {JsonKey("/rare", JSONType::STRING), JsonKeyLayoutType::SHARED}, + }; + std::vector field_datas = { + CreateMaterializeJsonFieldData({}), + CreateMaterializeJsonFieldData({"", "{}", " "}), + CreateNullableMaterializeJsonFieldData( + { + R"({"int":99,"text":"ignored","rare":"ignored"})", + "", + "{}", + R"({"rare":"also-ignored"})", + }, + {false, true, true, false}), + CreateMaterializeJsonFieldData({}), + }; + auto legacy_ranges = CreateJsonStatsRowRanges(field_datas, 2); + auto legacy = BuildMaterializeSnapshot( + legacy_ranges, key_types, {"/rare"}, true, false, 3); + ASSERT_NE(legacy.parquet.table, nullptr); + EXPECT_EQ(legacy.parquet.table->num_rows(), 7); + EXPECT_TRUE(legacy.postings.empty()); + + struct Scenario { + int64_t range_rows; + int64_t write_batch_rows; + int worker_count; + }; + const std::vector scenarios = { + {1, 1, 1}, + {3, 2, 2}, + {64, 16, 4}, + }; + + auto original_thread_num = + milvus::futures::getJsonStatsBuildExecutor()->numThreads(); + auto restore_threads = folly::makeGuard([&]() { + executor_set_json_stats_build_thread_num( + static_cast(original_thread_num)); + }); + for (const auto& scenario : scenarios) { + SCOPED_TRACE(::testing::Message() + << "range_rows=" << scenario.range_rows + << ", write_batch_rows=" << scenario.write_batch_rows + << ", worker_count=" << scenario.worker_count); + executor_set_json_stats_build_thread_num(scenario.worker_count); + auto row_ranges = + CreateJsonStatsRowRanges(field_datas, scenario.range_rows); + auto parallel = BuildMaterializeSnapshot(row_ranges, + key_types, + {"/rare"}, + true, + true, + scenario.write_batch_rows); + ExpectMaterializeSnapshotsEqual(legacy, parallel); + } +} + +TEST(JsonStatsMaterializeTest, + ParallelMatchesLegacyAcrossRangeAndWriterBatchMatrix) { + auto dataset = CreateMixedMaterializeDataset(); + auto legacy_ranges = CreateJsonStatsRowRanges(dataset.field_datas, 9); + auto legacy = BuildMaterializeSnapshot( + legacy_ranges, dataset.key_types, dataset.shared_paths, true, false, 5); + ASSERT_NE(legacy.parquet.table, nullptr); + EXPECT_EQ(legacy.parquet.table->num_rows(), dataset.row_count); + ASSERT_TRUE(legacy.postings.find("/dense") != legacy.postings.end()); + const auto& dense_postings = legacy.postings.at("/dense"); + ASSERT_EQ(dense_postings.size(), dataset.expected_dense_postings); + for (size_t i = 1; i < dense_postings.size(); ++i) { + EXPECT_LT(dense_postings[i - 1].first, dense_postings[i].first); + } + + struct Scenario { + int64_t range_rows; + int64_t write_batch_rows; + int worker_count; + }; + const std::vector scenarios = { + {1, 1, 4}, + {2, 17, 2}, + {7, 3, 4}, + {16, 64, 3}, + {128, 5, 4}, + }; + + auto original_thread_num = + milvus::futures::getJsonStatsBuildExecutor()->numThreads(); + auto restore_threads = folly::makeGuard([&]() { + executor_set_json_stats_build_thread_num( + static_cast(original_thread_num)); + }); + for (const auto& scenario : scenarios) { + SCOPED_TRACE(::testing::Message() + << "range_rows=" << scenario.range_rows + << ", write_batch_rows=" << scenario.write_batch_rows + << ", worker_count=" << scenario.worker_count); + executor_set_json_stats_build_thread_num(scenario.worker_count); + auto row_ranges = + CreateJsonStatsRowRanges(dataset.field_datas, scenario.range_rows); + auto parallel = BuildMaterializeSnapshot(row_ranges, + dataset.key_types, + dataset.shared_paths, + true, + true, + scenario.write_batch_rows); + ExpectMaterializeSnapshotsEqual(legacy, parallel); + } +} + +TEST(JsonStatsMaterializeTest, ParallelMatchesLegacyWithSharedOnlySchema) { + auto dataset = CreateMixedMaterializeDataset(); + for (auto& entry : dataset.key_types) { + entry.second = JsonKeyLayoutType::SHARED; + } + dataset.shared_paths = { + "/id", + "/flag", + "/text", + "/score", + "/arr", + "/mixed", + "/dense", + "/nested", + "/nested/value", + "/a~1b", + "/a~1b/~0key", + "/empty_obj", + "/nullv", + }; + + auto legacy = BuildMaterializeSnapshot( + CreateJsonStatsRowRanges(dataset.field_datas, 6), + dataset.key_types, + dataset.shared_paths, + true, + false, + 5); + ASSERT_NE(legacy.parquet.table, nullptr); + ASSERT_EQ(legacy.parquet.table->num_columns(), 1); + ASSERT_EQ(legacy.postings.at("/dense").size(), + dataset.expected_dense_postings); + + struct Scenario { + int64_t range_rows; + int64_t write_batch_rows; + int worker_count; + }; + const std::vector scenarios = { + {1, 1, 1}, + {7, 3, 4}, + {128, 64, 4}, + }; + + auto original_thread_num = + milvus::futures::getJsonStatsBuildExecutor()->numThreads(); + auto restore_threads = folly::makeGuard([&]() { + executor_set_json_stats_build_thread_num( + static_cast(original_thread_num)); + }); + for (const auto& scenario : scenarios) { + SCOPED_TRACE(::testing::Message() + << "range_rows=" << scenario.range_rows + << ", write_batch_rows=" << scenario.write_batch_rows + << ", worker_count=" << scenario.worker_count); + executor_set_json_stats_build_thread_num(scenario.worker_count); + auto parallel = BuildMaterializeSnapshot( + CreateJsonStatsRowRanges(dataset.field_datas, scenario.range_rows), + dataset.key_types, + dataset.shared_paths, + true, + true, + scenario.write_batch_rows); + ASSERT_EQ(parallel.parquet.table->num_columns(), 1); + ExpectMaterializeSnapshotsEqual(legacy, parallel); + } +} + +TEST(JsonStatsMaterializeTest, + TinyMemoryBudgetMatchesLegacyForMultipleAndOversizedChunks) { + auto legacy_dataset = CreateMixedMaterializeDataset(); + auto parallel_dataset = CreateMixedMaterializeDataset(); + auto legacy = BuildMaterializeSnapshot( + CreateJsonStatsRowRanges(legacy_dataset.field_datas, 8), + legacy_dataset.key_types, + legacy_dataset.shared_paths, + true, + false, + 5); + + auto original_thread_num = + milvus::futures::getJsonStatsBuildExecutor()->numThreads(); + auto& budget = JsonStatsBuildMemoryBudget::GetInstance(); + ASSERT_EQ(budget.InflightBytes(), 0); + auto original_capacity = budget.CapacityBytes(); + auto restore_resources = folly::makeGuard([&]() { + executor_set_json_stats_build_thread_num( + static_cast(original_thread_num)); + executor_set_json_stats_build_max_inflight_bytes( + static_cast(original_capacity)); + }); + + executor_set_json_stats_build_thread_num(4); + executor_set_json_stats_build_max_inflight_bytes(1); + + auto many_chunks = BuildMaterializeSnapshot( + CreateJsonStatsRowRanges(parallel_dataset.field_datas, 1), + parallel_dataset.key_types, + parallel_dataset.shared_paths, + true, + true, + 2); + ExpectMaterializeSnapshotsEqual(legacy, many_chunks); + EXPECT_EQ(budget.InflightBytes(), 0); + + auto oversized_dataset = CreateMixedMaterializeDataset(); + auto one_oversized_chunk = BuildMaterializeSnapshot( + CreateJsonStatsRowRanges(oversized_dataset.field_datas, + oversized_dataset.row_count), + oversized_dataset.key_types, + oversized_dataset.shared_paths, + true, + true, + 64); + ExpectMaterializeSnapshotsEqual(legacy, one_oversized_chunk); + EXPECT_EQ(budget.InflightBytes(), 0); +} + +TEST(JsonStatsMaterializeTest, ConcurrentBuildsSharingExecutorMatchLegacy) { + auto legacy_dataset = CreateMixedMaterializeDataset(); + auto first_dataset = CreateMixedMaterializeDataset(); + auto second_dataset = CreateMixedMaterializeDataset(); + auto legacy = BuildMaterializeSnapshot( + CreateJsonStatsRowRanges(legacy_dataset.field_datas, 8), + legacy_dataset.key_types, + legacy_dataset.shared_paths, + true, + false, + 5); + + auto original_thread_num = + milvus::futures::getJsonStatsBuildExecutor()->numThreads(); + auto& budget = JsonStatsBuildMemoryBudget::GetInstance(); + ASSERT_EQ(budget.InflightBytes(), 0); + auto original_capacity = budget.CapacityBytes(); + auto restore_resources = folly::makeGuard([&]() { + executor_set_json_stats_build_thread_num( + static_cast(original_thread_num)); + executor_set_json_stats_build_max_inflight_bytes( + static_cast(original_capacity)); + }); + executor_set_json_stats_build_thread_num(4); + executor_set_json_stats_build_max_inflight_bytes(1); + + auto first_future = std::async(std::launch::async, [&]() { + return BuildMaterializeSnapshot( + CreateJsonStatsRowRanges(first_dataset.field_datas, 1), + first_dataset.key_types, + first_dataset.shared_paths, + true, + true, + 2); + }); + auto second_future = std::async(std::launch::async, [&]() { + return BuildMaterializeSnapshot( + CreateJsonStatsRowRanges(second_dataset.field_datas, 13), + second_dataset.key_types, + second_dataset.shared_paths, + true, + true, + 11); + }); + + ASSERT_EQ(first_future.wait_for(std::chrono::seconds(30)), + std::future_status::ready); + ASSERT_EQ(second_future.wait_for(std::chrono::seconds(30)), + std::future_status::ready); + auto first = first_future.get(); + auto second = second_future.get(); + ExpectMaterializeSnapshotsEqual(legacy, first); + ExpectMaterializeSnapshotsEqual(legacy, second); + ExpectMaterializeSnapshotsEqual(first, second); + EXPECT_EQ(budget.InflightBytes(), 0); +} + +TEST(JsonStatsMaterializeTest, + EightConcurrentBuildsWithTinyBudgetAndFailuresComplete) { + constexpr int kWorkerCount = 8; + constexpr int kConcurrentBuildCount = 8; + constexpr int kStressRounds = 20; + constexpr int kRowCount = 33; + constexpr auto kRoundTimeout = std::chrono::seconds(30); + + const std::map key_types = { + {JsonKey("/int", JSONType::INT64), JsonKeyLayoutType::TYPED}, + }; + + auto original_thread_num = + milvus::futures::getJsonStatsBuildExecutor()->numThreads(); + auto& budget = JsonStatsBuildMemoryBudget::GetInstance(); + ASSERT_EQ(budget.InflightBytes(), 0); + auto original_capacity = budget.CapacityBytes(); + auto restore_resources = folly::makeGuard([&]() { + executor_set_json_stats_build_thread_num( + static_cast(original_thread_num)); + executor_set_json_stats_build_max_inflight_bytes( + static_cast(original_capacity)); + }); + executor_set_json_stats_build_thread_num(kWorkerCount); + executor_set_json_stats_build_max_inflight_bytes(1); + + for (int round = 0; round < kStressRounds; ++round) { + SCOPED_TRACE(::testing::Message() << "round=" << round); + std::barrier<> start_barrier(kConcurrentBuildCount + 1); + std::vector> build_futures; + std::vector expect_failure; + build_futures.reserve(kConcurrentBuildCount); + expect_failure.reserve(kConcurrentBuildCount); + + for (int build = 0; build < kConcurrentBuildCount; ++build) { + std::vector rows; + rows.reserve(kRowCount); + for (int row = 0; row < kRowCount; ++row) { + rows.push_back(fmt::format(R"({{"int":{}}})", row)); + } + + const bool inject_failure = build % 2 != 0; + if (inject_failure) { + // Rotate failures across the first, middle, and last ranges so + // exception cleanup is exercised after different amounts of + // budget handoff and completed work. + const int failure_slot = (round + build / 2) % 3; + const int failure_row = failure_slot == 0 ? 0 + : failure_slot == 1 ? kRowCount / 2 + : kRowCount - 1; + rows[failure_row] = R"({"int":)"; + } + expect_failure.push_back(inject_failure); + + std::vector field_datas = { + CreateMaterializeJsonFieldData(rows), + }; + auto row_ranges = + CreateJsonStatsRowRanges(field_datas, 1 + (round + build) % 7); + build_futures.emplace_back(std::async( + std::launch::async, + [&start_barrier, + row_ranges = std::move(row_ranges), + key_types, + write_batch_size = 1 + build % 5]() mutable { + start_barrier.arrive_and_wait(); + auto test_index = + CreateMaterializeTestIndex(false, write_batch_size); + BuildKeyStatsAccessor::Prepare( + *test_index.stats, key_types, kRowCount); + try { + BuildKeyStatsAccessor::CallParallel( + *test_index.stats, row_ranges, false); + BuildKeyStatsAccessor::Finish(*test_index.stats); + return std::string{}; + } catch (const std::exception& e) { + return std::string(e.what()); + } catch (...) { + return std::string("non-standard exception"); + } + })); + } + + start_barrier.arrive_and_wait(); + const auto deadline = std::chrono::steady_clock::now() + kRoundTimeout; + for (int build = 0; build < kConcurrentBuildCount; ++build) { + ASSERT_EQ(build_futures[build].wait_until(deadline), + std::future_status::ready) + << "concurrent json stats build timed out: round=" << round + << ", build=" << build; + } + + for (int build = 0; build < kConcurrentBuildCount; ++build) { + auto error = build_futures[build].get(); + if (expect_failure[build]) { + EXPECT_NE(error.find("Failed to parse Json"), std::string::npos) + << "expected injected parse failure: round=" << round + << ", build=" << build << ", error=" << error; + } else { + EXPECT_TRUE(error.empty()) + << "valid concurrent build failed: round=" << round + << ", build=" << build << ", error=" << error; + } + } + EXPECT_EQ(budget.InflightBytes(), 0) + << "transient budget leaked after round " << round; + } +} + +TEST(JsonStatsMaterializeTest, ParallelRethrowsLegacyParseError) { + std::map key_types = { + {JsonKey("/int", JSONType::INT64), JsonKeyLayoutType::TYPED}, + }; + auto original_thread_num = + milvus::futures::getJsonStatsBuildExecutor()->numThreads(); + auto& budget = JsonStatsBuildMemoryBudget::GetInstance(); + ASSERT_EQ(budget.InflightBytes(), 0); + auto original_capacity = budget.CapacityBytes(); + auto restore_resources = folly::makeGuard([&]() { + executor_set_json_stats_build_thread_num( + static_cast(original_thread_num)); + executor_set_json_stats_build_max_inflight_bytes( + static_cast(original_capacity)); + }); + executor_set_json_stats_build_max_inflight_bytes(1); + + for (int error_row : {0, 4, 8}) { + SCOPED_TRACE(::testing::Message() << "error_row=" << error_row); + std::vector rows; + for (int row = 0; row < 9; ++row) { + rows.push_back(fmt::format(R"({{"int":{}}})", row)); + } + rows[error_row] = R"({"int":)"; + std::vector field_datas = { + CreateMaterializeJsonFieldData(rows), + }; + auto row_ranges = CreateJsonStatsRowRanges(field_datas, 2); + ASSERT_EQ(row_ranges.size(), 5); + + auto capture_error = [&](bool parallel) { + auto test_index = CreateMaterializeTestIndex(false, 3); + BuildKeyStatsAccessor::Prepare(*test_index.stats, key_types, 9); + try { + if (parallel) { + BuildKeyStatsAccessor::CallParallel( + *test_index.stats, row_ranges, false); + } else { + BuildKeyStatsAccessor::CallLegacy( + *test_index.stats, row_ranges, false); + } + } catch (const std::exception& e) { + return std::string(e.what()); + } + return std::string{}; + }; + + auto legacy_error = capture_error(false); + ASSERT_FALSE(legacy_error.empty()); + EXPECT_NE(legacy_error.find("Failed to parse Json"), std::string::npos); + for (int worker_count : {1, 2, 4}) { + SCOPED_TRACE(::testing::Message() + << "worker_count=" << worker_count); + executor_set_json_stats_build_thread_num(worker_count); + auto parallel_error = capture_error(true); + ASSERT_FALSE(parallel_error.empty()); + EXPECT_EQ(parallel_error, legacy_error); + EXPECT_EQ(budget.InflightBytes(), 0); + } + } + + std::vector valid_field_datas = { + CreateMaterializeJsonFieldData( + {R"({"int":1})", R"({"int":2})", R"({"int":3})"}), + }; + auto valid_ranges = CreateJsonStatsRowRanges(valid_field_datas, 1); + auto valid_legacy = + BuildMaterializeSnapshot(valid_ranges, key_types, {}, false, false, 2); + auto valid_parallel = + BuildMaterializeSnapshot(valid_ranges, key_types, {}, false, true, 2); + ExpectMaterializeSnapshotsEqual(valid_legacy, valid_parallel); + EXPECT_EQ(budget.InflightBytes(), 0); +} + class JsonKeyStatsTest : public ::testing::TestWithParam { protected: void diff --git a/internal/core/unittest/test_json_stats/test_parquet_writer.cpp b/internal/core/unittest/test_json_stats/test_parquet_writer.cpp index 2b901c1278..d2497c44aa 100644 --- a/internal/core/unittest/test_json_stats/test_parquet_writer.cpp +++ b/internal/core/unittest/test_json_stats/test_parquet_writer.cpp @@ -1,8 +1,10 @@ #include +#include #include #include #include #include +#include #include #include #include @@ -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(&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(&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(&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(); + 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 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(input_builders.front()); + auto shared_builder = + std::static_pointer_cast(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> arrays; + arrays.reserve(input_builders.size()); + for (auto& builder : input_builders) { + std::shared_ptr 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 diff --git a/internal/core/unittest/test_json_stats/test_traverse_json_for_build_stats.cpp b/internal/core/unittest/test_json_stats/test_traverse_json_for_build_stats.cpp index c8eea5502e..b97840f4e7 100644 --- a/internal/core/unittest/test_json_stats/test_traverse_json_for_build_stats.cpp +++ b/internal/core/unittest/test_json_stats/test_traverse_json_for_build_stats.cpp @@ -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 +#include #include #include +#include #include #include +#include #include #include @@ -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 + CallLegacy(JsonKeyStats& stats, + const std::vector& field_datas, + bool nullable) { + return stats.CollectKeyInfo(field_datas, nullable); + } + + static std::map + CallParallel( + JsonKeyStats& stats, + const std::vector& 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 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& json_rows) { + auto field_data = milvus::storage::CreateFieldData( + milvus::DataType::JSON, milvus::DataType::NONE, false); + if (json_rows.empty()) { + return field_data; + } + + std::vector 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& json_rows, + const std::vector& 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 rows; + rows.reserve(json_rows.size()); + for (const auto& json : json_rows) { + rows.emplace_back(simdjson::padded_string(json)); + } + + std::vector 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(1U << (i % 8)); + } + } + field_data->FillFieldData(rows.data(), valid_bitmap.data(), rows.size(), 0); + return field_data; +} + +std::unique_ptr +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(ctx, true); +} + +void +ExpectKeyInfoMapsEqual( + const std::map& expected, + const std::map& 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& 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 +ExpectParallelCollectMatchesLegacy( + const std::vector& field_datas, + bool nullable, + const std::vector& rows_per_range_values, + const std::vector& 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(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 field_datas = { + CreateJsonFieldData(0), + CreateJsonFieldData(2), + CreateJsonFieldData(0), + CreateJsonFieldData(5), + CreateJsonFieldData(1), + CreateJsonFieldData(0), + }; + + using RowPosition = + std::tuple; + std::vector 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 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{ + R"({"a":1,"b":"x","nested":{"flag":true}})", + R"({"a":2,"b":"y"})", + R"({"a":"mixed","items":[1,2]})", + }); + auto second = CreateJsonFieldData(std::vector{ + R"({})", + R"({"nested":{"flag":false},"null_value":null})", + R"({"a":3,"b":"z"})", + R"([])", + "", + }); + std::vector 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::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 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{ + "42", + "-42", + "0", + "1.5", + "6.02e23", + "true", + "false", + "null", + "\"root string\"", + "[]", + R"([1,{"array_child":"is_not_traversed"}])", + "{}", + }); + auto second = CreateJsonFieldData(std::vector{ + 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 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{R"({"always":true})"}); + std::vector 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 field_datas = { + CreateJsonFieldData(0), + CreateJsonFieldData(std::vector{R"({"value":1})"}), + CreateJsonFieldData(0), + CreateJsonFieldData(std::vector{ + 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{ + 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 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{R"({"broken":)"}); + auto valid = CreateJsonFieldData( + std::vector{R"({"value":1})", R"({"value":2})"}); + std::vector 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"( diff --git a/internal/core/unittest/test_json_stats/test_utils.cpp b/internal/core/unittest/test_json_stats/test_utils.cpp index 0f6ec03379..f18e23415a 100644 --- a/internal/core/unittest/test_json_stats/test_utils.cpp +++ b/internal/core/unittest/test_json_stats/test_utils.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -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 column_map = { + {JsonKey("/typed", JSONType::INT64), JsonKeyLayoutType::TYPED}, + {JsonKey("/dynamic", JSONType::STRING), JsonKeyLayoutType::DYNAMIC}, + }; + std::set 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 column_map = { {JsonKey("int_key", JSONType::INT64), JsonKeyLayoutType::TYPED}, @@ -229,4 +272,4 @@ TEST_F(JsonStatsMetaTest, EmptyLayoutTypeMapTest) { EXPECT_TRUE(key_field_map.empty()); } -} // namespace milvus::index \ No newline at end of file +} // namespace milvus::index diff --git a/internal/datanode/index/init_segcore.go b/internal/datanode/index/init_segcore.go index 51328000b1..1b3f393edd 100644 --- a/internal/datanode/index/init_segcore.go +++ b/internal/datanode/index/init_segcore.go @@ -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() diff --git a/internal/util/initcore/json_stats_executor.go b/internal/util/initcore/json_stats_executor.go new file mode 100644 index 0000000000..25d4e72eff --- /dev/null +++ b/internal/util/initcore/json_stats_executor.go @@ -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)) +} diff --git a/internal/util/initcore/json_stats_executor_test.go b/internal/util/initcore/json_stats_executor_test.go new file mode 100644 index 0000000000..3b88b59094 --- /dev/null +++ b/internal/util/initcore/json_stats_executor_test.go @@ -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()) +} diff --git a/pkg/util/paramtable/component_param.go b/pkg/util/paramtable/component_param.go index df3e28a70d..867e2fd8d9 100644 --- a/pkg/util/paramtable/component_param.go +++ b/pkg/util/paramtable/component_param.go @@ -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", diff --git a/pkg/util/paramtable/component_param_test.go b/pkg/util/paramtable/component_param_test.go index 3eea5437f0..7b7a4d70f6 100644 --- a/pkg/util/paramtable/component_param_test.go +++ b/pkg/util/paramtable/component_param_test.go @@ -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") diff --git a/scripts/segcore_error_boundary_baseline.txt b/scripts/segcore_error_boundary_baseline.txt index a6374add45..a93a07d57d 100644 --- a/scripts/segcore_error_boundary_baseline.txt +++ b/scripts/segcore_error_boundary_baseline.txt @@ -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