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