diff --git a/internal/core/src/query/generated/ExecPlanNodeVisitor.h b/internal/core/src/query/generated/ExecPlanNodeVisitor.h index 1c0b47abb0..b055330d64 100644 --- a/internal/core/src/query/generated/ExecPlanNodeVisitor.h +++ b/internal/core/src/query/generated/ExecPlanNodeVisitor.h @@ -60,6 +60,7 @@ class ExecPlanNodeVisitor : public PlanNodeVisitor { node.accept(*this); assert(retrieve_ret_.has_value()); auto retrieve_ret = std::move(retrieve_ret_).value(); + retrieve_ret_.reset(); retrieve_ret_ = std::nullopt; return retrieve_ret; } diff --git a/internal/core/src/query/visitors/ExecPlanNodeVisitor.cpp b/internal/core/src/query/visitors/ExecPlanNodeVisitor.cpp index 0c30f70428..7f25aa4058 100644 --- a/internal/core/src/query/visitors/ExecPlanNodeVisitor.cpp +++ b/internal/core/src/query/visitors/ExecPlanNodeVisitor.cpp @@ -125,6 +125,7 @@ ExecPlanNodeVisitor::visit(RetrievePlanNode& node) { auto active_count = segment->get_active_count(timestamp_); if (active_count == 0) { + retrieve_ret_ = ret; return; } diff --git a/internal/core/src/segcore/SegmentSealedImpl.cpp b/internal/core/src/segcore/SegmentSealedImpl.cpp index 685edbf2b2..d57097c410 100644 --- a/internal/core/src/segcore/SegmentSealedImpl.cpp +++ b/internal/core/src/segcore/SegmentSealedImpl.cpp @@ -499,6 +499,10 @@ SegmentSealedImpl::mask_with_timestamps(boost::dynamic_bitset<>& bitset_chunk, T // just skip return; } + if (range.first == range.second && range.first == 0) { + bitset_chunk.reset(); + return; + } auto mask = TimestampIndex::GenerateBitset(timestamp, range, this->timestamps_.data(), this->timestamps_.size()); bitset_chunk &= mask; } diff --git a/internal/core/unittest/test_get_entity_by_ids.cpp b/internal/core/unittest/test_get_entity_by_ids.cpp index 3796a0b20f..8169510adf 100644 --- a/internal/core/unittest/test_get_entity_by_ids.cpp +++ b/internal/core/unittest/test_get_entity_by_ids.cpp @@ -152,6 +152,43 @@ TEST(Retrieve, AUTOID) { ASSERT_EQ(field1_data.data_size(), DIM * req_size); } +TEST(Retrieve2, LargeTimestamp) { + auto schema = std::make_shared(); + auto fid_64 = schema->AddDebugField("i64", DataType::INT64); + auto DIM = 16; + auto fid_vec = schema->AddDebugField("vector_64", DataType::VECTOR_FLOAT, DIM, MetricType::METRIC_L2); + schema->set_primary_key(FieldOffset(0)); + + int64_t N = 100; + int64_t req_size = 10; + auto choose = [=](int i) { return i * 3 % N; }; + uint64_t ts_offset = 100; + auto dataset = DataGen(schema, N, 42, ts_offset + 1); + auto segment = CreateSealedSegment(schema); + SealedLoader(dataset, *segment); + auto i64_col = dataset.get_col(0); + + auto plan = std::make_unique(*schema); + + auto term_expr = std::make_unique>(); + term_expr->field_offset_ = FieldOffset(0); + term_expr->data_type_ = DataType::INT64; + for (int i = 0; i < req_size; ++i) { + term_expr->terms_.emplace_back(i64_col[choose(i)]); + } + plan->plan_node_ = std::make_unique(); + plan->plan_node_->predicate_ = std::move(term_expr); + std::vector target_offsets{FieldOffset(0), FieldOffset(1)}; + plan->field_offsets_ = target_offsets; + + auto retrieve_results = segment->Retrieve(plan.get(), ts_offset); + Assert(retrieve_results->fields_data_size() == 2); + auto field0 = retrieve_results->fields_data(0); + auto field1 = retrieve_results->fields_data(1); + Assert(field0.scalars().long_data().data_size() == 0); + Assert(field1.scalars().long_data().data_size() == 0); +} + TEST(GetEntityByIds, PrimaryKey) { auto schema = std::make_shared(); auto fid_64 = schema->AddDebugField("counter_i64", DataType::INT64); diff --git a/internal/core/unittest/test_utils/DataGen.h b/internal/core/unittest/test_utils/DataGen.h index 95c9e6ed4d..1b001965c2 100644 --- a/internal/core/unittest/test_utils/DataGen.h +++ b/internal/core/unittest/test_utils/DataGen.h @@ -55,7 +55,7 @@ struct GeneratedData { private: GeneratedData() = default; friend GeneratedData - DataGen(SchemaPtr schema, int64_t N, uint64_t seed); + DataGen(SchemaPtr schema, int64_t N, uint64_t seed, uint64_t ts_offset); void generate_rows(int64_t N, SchemaPtr schema); }; @@ -86,7 +86,7 @@ GeneratedData::generate_rows(int64_t N, SchemaPtr schema) { } inline GeneratedData -DataGen(SchemaPtr schema, int64_t N, uint64_t seed = 42) { +DataGen(SchemaPtr schema, int64_t N, uint64_t seed = 42, uint64_t ts_offset = 0) { using std::vector; std::vector> cols; std::default_random_engine er(seed); @@ -193,7 +193,7 @@ DataGen(SchemaPtr schema, int64_t N, uint64_t seed = 42) { res.cols_ = std::move(cols); for (int i = 0; i < N; ++i) { res.row_ids_.push_back(i); - res.timestamps_.push_back(i); + res.timestamps_.push_back(i + ts_offset); } // std::shuffle(res.row_ids_.begin(), res.row_ids_.end(), er); res.generate_rows(N, schema);