From 1de4d8a7f16256f176a9d22ac888f4a864dac302 Mon Sep 17 00:00:00 2001 From: sijie-ni-0214 Date: Thu, 10 Sep 2026 09:11:37 +0000 Subject: [PATCH] enhance: parallelize retrieve output field materialization Share field-level scheduling between search and retrieve. Run field reads on the middle-priority pool with independent operation contexts, then publish results in plan order while preserving primary key handling and storage accounting. Cancel sibling reads on failure, preserve the original exception, and drain pending tasks before returning. Reserve result storage before dispatch and add coverage for concurrent reads, cancellation, failure cleanup, and schema-dependent output fields. Signed-off-by: sijie-ni-0214 --- .../core/src/segcore/SegmentInterface.cpp | 248 ++++++----- .../SegmentRetrieveOutputFieldsTest.cpp | 395 ++++++++++++++++++ 2 files changed, 529 insertions(+), 114 deletions(-) create mode 100644 internal/core/src/segcore/SegmentRetrieveOutputFieldsTest.cpp diff --git a/internal/core/src/segcore/SegmentInterface.cpp b/internal/core/src/segcore/SegmentInterface.cpp index 30b472b2ee3..c14690022d4 100644 --- a/internal/core/src/segcore/SegmentInterface.cpp +++ b/internal/core/src/segcore/SegmentInterface.cpp @@ -68,6 +68,97 @@ struct FetchedOutputField { int64_t scanned_total_bytes; }; +template +std::vector +FetchOutputFields(const std::vector& field_ids, + int64_t segment_id, + milvus::OpContext* op_ctx, + FetchField fetch_field) { + if (field_ids.empty()) { + return {}; + } + + // Reserve both containers before starting workers. No allocation while + // collecting results may bypass draining tasks that borrow caller state. + std::vector> futures; + futures.reserve(field_ids.size()); + std::vector fetched_fields; + fetched_fields.reserve(field_ids.size()); + + folly::CancellationSource sibling_cancel_source; + const auto field_cancellation_token = + op_ctx != nullptr + ? folly::cancellation_token_merge(op_ctx->cancellation_token, + sibling_cancel_source.getToken()) + : sibling_cancel_source.getToken(); + std::mutex field_error_mutex; + std::exception_ptr first_field_error; + auto record_field_error = [&]() { + std::lock_guard lock(field_error_mutex); + if (first_field_error == nullptr) { + first_field_error = std::current_exception(); + } + }; + + // Search and retrieve parent tasks run on the search executor. Keep + // field tasks on MIDDLE so parents cannot starve their own children. + auto& pool = ThreadPools::GetThreadPool(ThreadPoolPriority::MIDDLE); + try { + for (auto field_id : field_ids) { + futures.emplace_back(pool.Submit([&, field_id]() { + try { + milvus::OpContext field_ctx; + field_ctx.cancellation_token = field_cancellation_token; + if (op_ctx != nullptr) { + field_ctx.runtime_load_priority = + op_ctx->runtime_load_priority; + field_ctx.coload_fields = op_ctx->coload_fields; + field_ctx.pinned_segment_state = + op_ctx->pinned_segment_state; + field_ctx.pinned_state_owner = + op_ctx->pinned_state_owner; + field_ctx.trace_context = op_ctx->trace_context; + field_ctx.trace_span = op_ctx->trace_span; + } + segcore::CheckCancellation(&field_ctx, + segment_id, + field_id.get(), + "FillTargetEntry"); + auto field_data = fetch_field(field_id, &field_ctx); + return FetchedOutputField{ + field_id, + std::move(field_data), + field_ctx.storage_usage.scanned_cold_bytes.load(), + field_ctx.storage_usage.scanned_total_bytes.load()}; + } catch (...) { + // Record the original error before cancelling siblings, + // so their cancellation cannot replace the root cause. + record_field_error(); + sibling_cancel_source.requestCancellation(); + throw; + } + })); + } + } catch (...) { + sibling_cancel_source.requestCancellation(); + storage::DrainFutures(futures); + throw; + } + + for (auto& future : futures) { + try { + fetched_fields.emplace_back(future.get()); + } catch (...) { + record_field_error(); + sibling_cancel_source.requestCancellation(); + } + } + if (first_field_error != nullptr) { + std::rethrow_exception(first_field_error); + } + return fetched_fields; +} + } // namespace std::shared_ptr @@ -144,42 +235,15 @@ SegmentInternalInterface::FillSearchResultOutputFields( const std::vector& field_ids, SearchResult& results, milvus::OpContext* op_ctx) const { - if (field_ids.empty()) { - return; - } - - folly::CancellationSource sibling_cancel_source; - const auto field_cancellation_token = - op_ctx != nullptr - ? folly::cancellation_token_merge(op_ctx->cancellation_token, - sibling_cancel_source.getToken()) - : sibling_cancel_source.getToken(); - const auto size = results.seg_offsets_.size(); - auto fetch_one = [this, - plan, - &results, - size, - op_ctx, - field_cancellation_token](FieldId field_id) { - milvus::OpContext field_ctx; - field_ctx.cancellation_token = field_cancellation_token; - if (op_ctx != nullptr) { - field_ctx.runtime_load_priority = op_ctx->runtime_load_priority; - field_ctx.coload_fields = op_ctx->coload_fields; - field_ctx.pinned_segment_state = op_ctx->pinned_segment_state; - field_ctx.pinned_state_owner = op_ctx->pinned_state_owner; - field_ctx.trace_context = op_ctx->trace_context; - field_ctx.trace_span = op_ctx->trace_span; - } - segcore::CheckCancellation( - &field_ctx, get_segment_id(), field_id.get(), "FillTargetEntry"); + auto fetch_one = [this, plan, &results, size]( + FieldId field_id, milvus::OpContext* field_ctx) { auto& field_meta = plan->schema_->operator[](field_id); std::unique_ptr field_data; if (plan->schema_->get_dynamic_field_id().has_value() && plan->schema_->get_dynamic_field_id().value() == field_id && !plan->target_dynamic_fields_.empty()) { - field_data = bulk_subscript(&field_ctx, + field_data = bulk_subscript(field_ctx, field_id, results.seg_offsets_.data(), size, @@ -188,65 +252,13 @@ SegmentInternalInterface::FillSearchResultOutputFields( field_data = bulk_subscript_not_exist_field(field_meta, size); } else { field_data = bulk_subscript( - &field_ctx, field_id, results.seg_offsets_.data(), size); - } - return FetchedOutputField{ - field_id, - std::move(field_data), - field_ctx.storage_usage.scanned_cold_bytes.load(), - field_ctx.storage_usage.scanned_total_bytes.load()}; - }; - - std::vector> futures; - futures.reserve(field_ids.size()); - std::mutex field_error_mutex; - std::exception_ptr first_field_error; - auto record_field_error = [&field_error_mutex, &first_field_error]() { - std::lock_guard lock(field_error_mutex); - if (first_field_error == nullptr) { - first_field_error = std::current_exception(); + field_ctx, field_id, results.seg_offsets_.data(), size); } + return field_data; }; - // The caller runs one parent task per segment on the search executor. - // Fan fields out to MIDDLE so a segment never waits for children queued - // behind it in the same pool. - auto& pool = ThreadPools::GetThreadPool(ThreadPoolPriority::MIDDLE); - try { - for (auto field_id : field_ids) { - futures.emplace_back(pool.Submit([fetch_one, - field_id, - sibling_cancel_source, - record_field_error]() mutable { - try { - return fetch_one(field_id); - } catch (...) { - record_field_error(); - sibling_cancel_source.requestCancellation(); - throw; - } - })); - } - } catch (...) { - sibling_cancel_source.requestCancellation(); - storage::DrainFutures(futures); - throw; - } - - std::vector fetched_fields; - try { - fetched_fields = storage::WaitAllFutures(std::move(futures)); - } catch (...) { - std::exception_ptr field_error; - { - std::lock_guard lock(field_error_mutex); - field_error = first_field_error; - } - if (field_error != nullptr) { - std::rethrow_exception(field_error); - } - throw; - } + auto fetched_fields = + FetchOutputFields(field_ids, get_segment_id(), op_ctx, fetch_one); // Workers only return isolated values. Publish to the shared SearchResult // after every future completes so the map and counters stay race-free. @@ -654,23 +666,25 @@ SegmentInternalInterface::FillTargetEntry( return pk_field_id.has_value() && pk_field_id.value() == field_id; }; - // Per-call OpContext keeps storage_usage scoped to this segment; - // sharing the caller's op_ctx across segments would double-count - // bytes. Inherit the caller's cancellation_token and load priority so - // in-loop cancellation still propagates. - milvus::OpContext local_ctx; - if (op_ctx != nullptr) { - local_ctx.cancellation_token = op_ctx->cancellation_token; - local_ctx.runtime_load_priority = op_ctx->runtime_load_priority; - } + std::vector field_ids; + field_ids.reserve(plan->field_ids_.size()); for (auto field_id : plan->field_ids_) { + // System fields are needed for reduce even in the PK-only phase. + if (SystemProperty::Instance().IsSystem(field_id) || !ignore_non_pk || + is_pk_field(field_id)) { + field_ids.push_back(field_id); + } + } + + auto fetch_one = [this, plan, offsets, size](FieldId field_id, + milvus::OpContext* field_ctx) { if (SystemProperty::Instance().IsSystem(field_id)) { auto system_type = SystemProperty::Instance().GetSystemFieldType(field_id); FixedVector output(size); bulk_subscript( - &local_ctx, system_type, offsets, size, output.data()); + field_ctx, system_type, offsets, size, output.data()); auto data_array = std::make_unique(); data_array->set_field_id(field_id.get()); @@ -680,35 +694,48 @@ SegmentInternalInterface::FillTargetEntry( auto data = reinterpret_cast(output.data()); auto obj = scalar_array->mutable_long_data(); obj->mutable_data()->Add(data, data + size); - fields_data->AddAllocated(data_array.release()); - continue; - } - - if (ignore_non_pk && !is_pk_field(field_id)) { - continue; + return data_array; } if (plan->schema_->get_dynamic_field_id().has_value() && plan->schema_->get_dynamic_field_id().value() == field_id && !plan->target_dynamic_fields_.empty()) { auto& target_dynamic_fields = plan->target_dynamic_fields_; - auto col = bulk_subscript( - &local_ctx, field_id, offsets, size, target_dynamic_fields); - fields_data->AddAllocated(col.release()); - continue; + return bulk_subscript( + field_ctx, field_id, offsets, size, target_dynamic_fields); } std::unique_ptr col; auto& field_meta = plan->schema_->operator[](field_id); if (!is_field_exist(field_id)) { col = bulk_subscript_not_exist_field(field_meta, size); } else { - col = bulk_subscript(&local_ctx, field_id, offsets, size); + col = bulk_subscript(field_ctx, field_id, offsets, size); } // todo(SpadeA): consider vector array? if (field_meta.get_data_type() == DataType::ARRAY) { col->mutable_scalars()->mutable_array_data()->set_element_type( proto::schema::DataType(field_meta.get_element_type())); } + return col; + }; + + auto fetched_fields = + FetchOutputFields(field_ids, get_segment_id(), op_ctx, fetch_one); + + // Workers never mutate the result proto. Publish fields in plan order + // and account each field's IO once, after every task has completed. + for (auto& fetched : fetched_fields) { + auto field_id = fetched.field_id; + auto& col = fetched.field_data; + results->set_scanned_remote_bytes(results->scanned_remote_bytes() + + fetched.scanned_remote_bytes); + results->set_scanned_total_bytes(results->scanned_total_bytes() + + fetched.scanned_total_bytes); + if (SystemProperty::Instance().IsSystem(field_id)) { + fields_data->AddAllocated(col.release()); + continue; + } + auto& field_meta = plan->schema_->operator[](field_id); if (fill_ids && is_pk_field(field_id)) { // fill_ids should be true when the first Retrieve was called. The reduce phase depends on the ids to do // merge-sort. @@ -746,13 +773,6 @@ SegmentInternalInterface::FillTargetEntry( fields_data->AddAllocated(col.release()); } } - // Add retrieve_storage_cost to results - results->set_scanned_remote_bytes( - results->scanned_remote_bytes() + - local_ctx.storage_usage.scanned_cold_bytes.load()); - results->set_scanned_total_bytes( - results->scanned_total_bytes() + - local_ctx.storage_usage.scanned_total_bytes.load()); } std::unique_ptr diff --git a/internal/core/src/segcore/SegmentRetrieveOutputFieldsTest.cpp b/internal/core/src/segcore/SegmentRetrieveOutputFieldsTest.cpp new file mode 100644 index 00000000000..7c5b0de763b --- /dev/null +++ b/internal/core/src/segcore/SegmentRetrieveOutputFieldsTest.cpp @@ -0,0 +1,395 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/Consts.h" +#include "common/EasyAssert.h" +#include "common/OpContext.h" +#include "common/Schema.h" +#include "common/Utils.h" +#include "query/PlanImpl.h" +#include "segcore/SegmentGrowingImpl.h" +#include "segcore/Utils.h" +#include "test_utils/DataGen.h" + +using namespace milvus; +using namespace milvus::segcore; +using namespace std::chrono_literals; + +namespace { + +class InstrumentedRetrieveSegment : public SegmentGrowingImpl { + public: + explicit InstrumentedRetrieveSegment(SchemaPtr schema) + : SegmentGrowingImpl(std::move(schema), + nullptr, + SegcoreConfig::default_config(), + 101) { + } + + using SegmentGrowingImpl::bulk_subscript; + + std::unique_ptr + bulk_subscript(milvus::OpContext* op_ctx, + FieldId field_id, + const int64_t* offsets, + int64_t count) const override { + if (before_fetch) { + before_fetch(field_id, op_ctx); + } + return SegmentGrowingImpl::bulk_subscript( + op_ctx, field_id, offsets, count); + } + + std::function before_fetch; +}; + +class RetrieveOutputFieldsTest : public ::testing::TestWithParam { + protected: + void + SetUp() override { + schema = std::make_shared(); + pk = schema->AddDebugField("pk", GetParam()); + first = schema->AddDebugField("first", DataType::INT64); + last = schema->AddDebugField("last", DataType::INT64); + schema->set_primary_field_id(pk); + segment = std::make_unique(schema); + auto data = DataGen(schema, 4); + auto reserved = segment->PreInsert(4); + segment->Insert(reserved, + 4, + data.row_ids_.data(), + data.timestamps_.data(), + data.raw_); + plan = std::make_unique(schema); + plan->field_ids_ = {first, pk, last}; + } + + std::unique_ptr + ExpectedField(FieldId field_id) { + return segment->SegmentGrowingImpl::bulk_subscript( + nullptr, field_id, offsets.data(), offsets.size()); + } + + void + CheckIDs(const proto::segcore::RetrieveResults& results) { + auto expected = ExpectedField(pk); + if (GetParam() == DataType::INT64) { + EXPECT_EQ(results.ids().int_id().SerializeAsString(), + expected->scalars().long_data().SerializeAsString()); + } else { + EXPECT_EQ(results.ids().str_id().SerializeAsString(), + expected->scalars().string_data().SerializeAsString()); + } + } + + SchemaPtr schema; + FieldId pk{0}; + FieldId first{0}; + FieldId last{0}; + std::unique_ptr segment; + std::unique_ptr plan; + const std::vector offsets{3, 1, 3}; +}; + +TEST_P(RetrieveOutputFieldsTest, ParallelReadsPreserveOrderIDsAndStorageCost) { + std::promise last_started; + auto last_ready = last_started.get_future(); + milvus::OpContext parent_ctx; + parent_ctx.runtime_load_priority = 1; + parent_ctx.coload_fields = {first.get(), last.get()}; + parent_ctx.storage_usage.scanned_cold_bytes = 1000; + parent_ctx.storage_usage.scanned_total_bytes = 2000; + segment->before_fetch = [&](FieldId field_id, milvus::OpContext* ctx) { + EXPECT_NE(ctx, &parent_ctx); + EXPECT_EQ(ctx->runtime_load_priority, parent_ctx.runtime_load_priority); + EXPECT_EQ(ctx->coload_fields, parent_ctx.coload_fields); + EXPECT_EQ(ctx->storage_usage.scanned_cold_bytes.load(), 0); + EXPECT_EQ(ctx->storage_usage.scanned_total_bytes.load(), 0); + ctx->storage_usage.scanned_cold_bytes += field_id.get(); + ctx->storage_usage.scanned_total_bytes += field_id.get() * 2; + if (field_id == first) { + // A serial field loop cannot reach last while first is waiting. + EXPECT_EQ(last_ready.wait_for(5s), std::future_status::ready); + } else if (field_id == last) { + last_started.set_value(); + } + }; + + auto results = std::make_unique(); + results->set_scanned_remote_bytes(7); + results->set_scanned_total_bytes(11); + segment->FillTargetEntry(nullptr, + plan.get(), + results, + offsets.data(), + offsets.size(), + false, + true, + &parent_ctx); + + ASSERT_EQ(results->fields_data_size(), plan->field_ids_.size()); + for (size_t i = 0; i < plan->field_ids_.size(); ++i) { + EXPECT_EQ(results->fields_data(i).SerializeAsString(), + ExpectedField(plan->field_ids_[i])->SerializeAsString()); + } + CheckIDs(*results); + auto bytes = first.get() + pk.get() + last.get(); + EXPECT_EQ(results->scanned_remote_bytes(), 7 + bytes); + EXPECT_EQ(results->scanned_total_bytes(), 11 + bytes * 2); + EXPECT_EQ(parent_ctx.storage_usage.scanned_cold_bytes.load(), 1000); + EXPECT_EQ(parent_ctx.storage_usage.scanned_total_bytes.load(), 2000); +} + +TEST_P(RetrieveOutputFieldsTest, PKOnlyKeepsTimestampAndSkipsOtherReads) { + plan->field_ids_ = {first, TimestampFieldID, pk, last}; + std::atomic calls{0}; + segment->before_fetch = [&](FieldId field_id, milvus::OpContext*) { + EXPECT_EQ(field_id, pk); + ++calls; + }; + auto results = std::make_unique(); + segment->FillTargetEntry(nullptr, + plan.get(), + results, + offsets.data(), + offsets.size(), + true, + true); + + EXPECT_EQ(calls.load(), 1); + CheckIDs(*results); + ASSERT_EQ(results->fields_data_size(), 1); + EXPECT_EQ(results->fields_data(0).field_id(), TimestampFieldID.get()); + const auto& timestamps = results->fields_data(0).scalars().long_data(); + ASSERT_EQ(timestamps.data_size(), offsets.size()); + FixedVector expected_timestamps(offsets.size()); + segment->SegmentGrowingImpl::bulk_subscript(nullptr, + SystemFieldType::Timestamp, + offsets.data(), + offsets.size(), + expected_timestamps.data()); + for (size_t i = 0; i < offsets.size(); ++i) { + EXPECT_EQ(timestamps.data(i), expected_timestamps[i]); + } +} + +TEST_P(RetrieveOutputFieldsTest, RetrieveByOffsetsUsesParallelFieldReads) { + std::promise last_started; + auto last_ready = last_started.get_future(); + segment->before_fetch = [&](FieldId field_id, milvus::OpContext*) { + if (field_id == first) { + EXPECT_EQ(last_ready.wait_for(5s), std::future_status::ready); + } else if (field_id == last) { + last_started.set_value(); + } + }; + auto results = segment->Retrieve(nullptr, + plan.get(), + offsets.data(), + offsets.size(), + folly::CancellationToken()); + ASSERT_EQ(results->fields_data_size(), plan->field_ids_.size()); + for (size_t i = 0; i < plan->field_ids_.size(); ++i) { + EXPECT_EQ(results->fields_data(i).SerializeAsString(), + ExpectedField(plan->field_ids_[i])->SerializeAsString()); + } + EXPECT_EQ(results->ids().id_field_case(), IdArray::ID_FIELD_NOT_SET); +} + +TEST_P(RetrieveOutputFieldsTest, FailureCancelsAndDrainsSiblingsBeforeReturn) { + plan->field_ids_ = {first, last}; + std::promise first_started; + auto first_ready = first_started.get_future(); + std::promise sibling_cancelled; + auto cancelled = sibling_cancelled.get_future(); + std::promise release_first; + auto released = release_first.get_future(); + std::atomic first_exited{false}; + segment->before_fetch = [&](FieldId field_id, milvus::OpContext* ctx) { + if (field_id == first) { + folly::CancellationCallback on_cancel(ctx->cancellation_token, [&] { + sibling_cancelled.set_value(); + }); + first_started.set_value(); + EXPECT_EQ(released.wait_for(5s), std::future_status::ready); + first_exited = true; + CheckCancellation(ctx, segment->get_segment_id(), "test read"); + } else { + EXPECT_EQ(first_ready.wait_for(5s), std::future_status::ready); + throw SegcoreError(ErrorCode::FileReadFailed, + "injected read failure"); + } + }; + auto results = std::make_unique(); + auto pending = std::async(std::launch::async, [&] { + segment->FillTargetEntry(nullptr, + plan.get(), + results, + offsets.data(), + offsets.size(), + false, + false); + }); + EXPECT_EQ(cancelled.wait_for(5s), std::future_status::ready); + EXPECT_EQ(pending.wait_for(0s), std::future_status::timeout); + release_first.set_value(); + try { + pending.get(); + FAIL() << "Expected the original field read error"; + } catch (const SegcoreError& error) { + EXPECT_EQ(error.get_error_code(), ErrorCode::FileReadFailed); + EXPECT_STREQ(error.what(), "injected read failure"); + } + EXPECT_TRUE(first_exited.load()); + EXPECT_EQ(results->fields_data_size(), 0); + EXPECT_EQ(results->scanned_remote_bytes(), 0); +} + +TEST_P(RetrieveOutputFieldsTest, PreCancelledRetrieveDoesNotReadFields) { + std::atomic calls{0}; + segment->before_fetch = [&](FieldId, milvus::OpContext*) { ++calls; }; + folly::CancellationSource source; + source.requestCancellation(); + try { + segment->Retrieve(nullptr, + plan.get(), + offsets.data(), + offsets.size(), + source.getToken()); + FAIL() << "Expected cancellation"; + } catch (const SegcoreError& error) { + EXPECT_EQ(error.get_error_code(), ErrorCode::FollyCancel); + } + EXPECT_EQ(calls.load(), 0); +} + +TEST_P(RetrieveOutputFieldsTest, CancellationReachesRunningField) { + plan->field_ids_ = {first}; + folly::CancellationSource source; + std::promise started; + auto ready = started.get_future(); + segment->before_fetch = [&](FieldId, milvus::OpContext* ctx) { + std::promise cancelled; + auto cancelled_ready = cancelled.get_future(); + folly::CancellationCallback on_cancel(ctx->cancellation_token, + [&] { cancelled.set_value(); }); + started.set_value(); + EXPECT_EQ(cancelled_ready.wait_for(5s), std::future_status::ready); + CheckCancellation(ctx, segment->get_segment_id(), "test read"); + }; + auto pending = std::async(std::launch::async, [&] { + return segment->Retrieve(nullptr, + plan.get(), + offsets.data(), + offsets.size(), + source.getToken()); + }); + EXPECT_EQ(ready.wait_for(5s), std::future_status::ready); + source.requestCancellation(); + try { + pending.get(); + FAIL() << "Expected cancellation during the field read"; + } catch (const SegcoreError& error) { + EXPECT_EQ(error.get_error_code(), ErrorCode::FollyCancel); + } +} + +TEST_P(RetrieveOutputFieldsTest, AllocationFailurePreservesExceptionType) { + segment->before_fetch = [&](FieldId field_id, milvus::OpContext*) { + if (field_id == last) { + throw std::bad_alloc(); + } + }; + EXPECT_THROW(segment->Retrieve(nullptr, + plan.get(), + offsets.data(), + offsets.size(), + folly::CancellationToken()), + std::bad_alloc); +} + +TEST_P(RetrieveOutputFieldsTest, DynamicProjectionAndMissingFields) { + auto local_schema = std::make_shared(); + auto local_pk = local_schema->AddDebugField("pk", GetParam()); + local_schema->set_primary_field_id(local_pk); + auto dynamic = local_schema->AddDebugField("$meta", DataType::JSON); + local_schema->set_dynamic_field_id(dynamic); + auto array = + local_schema->AddDebugArrayField("array", DataType::INT64, false); + auto local_segment = CreateGrowingSegment(local_schema, nullptr); + auto data = DataGen(local_schema, 4); + for (auto& field : *data.raw_->mutable_fields_data()) { + if (field.field_id() == dynamic.get()) { + for (auto& json : *field.mutable_scalars() + ->mutable_json_data() + ->mutable_data()) { + json = R"({"keep":42,"drop":99})"; + } + } + } + auto reserved = local_segment->PreInsert(4); + local_segment->Insert( + reserved, 4, data.row_ids_.data(), data.timestamps_.data(), data.raw_); + + auto evolved = std::make_shared(*local_schema); + DefaultValueType default_value; + default_value.set_long_data(42); + auto added = evolved->AddDebugFieldWithDefaultValue( + "added", DataType::INT64, default_value); + auto missing_array = + evolved->AddDebugArrayField("missing_array", DataType::VARCHAR, true); + query::RetrievePlan local_plan(evolved); + local_plan.field_ids_ = {dynamic, added, array, missing_array, local_pk}; + local_plan.target_dynamic_fields_ = {"keep"}; + auto results = local_segment->Retrieve(nullptr, + &local_plan, + offsets.data(), + offsets.size(), + folly::CancellationToken()); + ASSERT_EQ(results->fields_data_size(), 5); + const auto& added_validity = + GetFieldDataRowValidData(results->fields_data(1)); + const auto& array_validity = + GetFieldDataRowValidData(results->fields_data(3)); + ASSERT_EQ(added_validity.size(), offsets.size()); + ASSERT_EQ(array_validity.size(), offsets.size()); + for (size_t i = 0; i < offsets.size(); ++i) { + EXPECT_EQ(results->fields_data(0).scalars().json_data().data(i), + R"({"keep":42})"); + EXPECT_EQ(results->fields_data(1).scalars().long_data().data(i), 42); + EXPECT_TRUE(added_validity[i]); + EXPECT_FALSE(array_validity[i]); + } + EXPECT_EQ(results->fields_data(2).scalars().array_data().element_type(), + proto::schema::DataType::Int64); + EXPECT_EQ(results->fields_data(3).scalars().array_data().element_type(), + proto::schema::DataType::VarChar); +} + +INSTANTIATE_TEST_SUITE_P(PrimaryKeyTypes, + RetrieveOutputFieldsTest, + ::testing::Values(DataType::INT64, DataType::VARCHAR)); + +} // namespace