From 947bff28ca571bb33703746eb5ba592cd029881b Mon Sep 17 00:00:00 2001 From: Spade A <71589810+SpadeA-Tang@users.noreply.github.com> Date: Tue, 23 Jun 2026 11:08:26 +0800 Subject: [PATCH] fix: support array length on struct array parent field (#50681) issue: https://github.com/milvus-io/milvus/issues/50510 ref: https://github.com/milvus-io/milvus/issues/42148 --------- Signed-off-by: SpadeA --- .../expression/BinaryArithOpEvalRangeExpr.cpp | 115 +++++++++++ .../expression/BinaryArithOpEvalRangeExpr.h | 9 +- internal/core/src/exec/expression/Expr.h | 190 ++++++++++++------ .../src/exec/expression/ExprArrayTest.cpp | 136 +++++++++++++ .../core/src/segcore/SegmentGrowingImpl.cpp | 86 +++++++- .../parser/planparserv2/parser_visitor.go | 17 +- .../planparserv2/plan_parser_v2_test.go | 74 ++++++- 7 files changed, 552 insertions(+), 75 deletions(-) diff --git a/internal/core/src/exec/expression/BinaryArithOpEvalRangeExpr.cpp b/internal/core/src/exec/expression/BinaryArithOpEvalRangeExpr.cpp index 60671818e7..fa23ffa046 100644 --- a/internal/core/src/exec/expression/BinaryArithOpEvalRangeExpr.cpp +++ b/internal/core/src/exec/expression/BinaryArithOpEvalRangeExpr.cpp @@ -24,6 +24,7 @@ #include "common/Array.h" #include "common/Json.h" #include "common/Tracer.h" +#include "common/VectorArray.h" #include "fmt/core.h" #include "opentelemetry/trace/span.h" @@ -119,6 +120,26 @@ PhyBinaryArithOpEvalRangeExpr::Eval(EvalCtx& context, VectorPtr& result) { } break; } + case DataType::VECTOR_ARRAY: { + auto value_type = expr_->value_.val_case(); + switch (value_type) { + case proto::plan::GenericValue::ValCase::kInt64Val: { + result = ExecRangeVisitorImplForVectorArray(input); + break; + } + case proto::plan::GenericValue::ValCase::kFloatVal: { + result = ExecRangeVisitorImplForVectorArray(input); + break; + } + default: { + ThrowInfo( + DataTypeInvalid, + fmt::format("unsupported value type {} in expression", + value_type)); + } + } + break; + } default: ThrowInfo(DataTypeInvalid, "unsupported data type: {}", @@ -922,6 +943,100 @@ PhyBinaryArithOpEvalRangeExpr::ExecRangeVisitorImplForArray( return res_vec; } +template +VectorPtr +PhyBinaryArithOpEvalRangeExpr::ExecRangeVisitorImplForVectorArray( + OffsetVector* input) { + if (expr_->arith_op_type_ != proto::plan::ArithOpType::ArrayLength) { + ThrowInfo(OpTypeInvalid, + "unsupported arith type for vector array field: {}", + expr_->arith_op_type_); + } + + auto real_batch_size = + has_offset_input_ ? input->size() : GetNextBatchSize(); + if (real_batch_size == 0) { + return nullptr; + } + + if (!arg_inited_) { + value_arg_.SetValue(expr_->value_); + arg_inited_ = true; + } + + auto res_vec = + std::make_shared(TargetBitmap(real_batch_size, false), + TargetBitmap(real_batch_size, true)); + TargetBitmapView res(res_vec->GetRawData(), real_batch_size); + TargetBitmapView valid_res(res_vec->GetValidRawData(), real_batch_size); + + auto op_type = expr_->op_type_; + auto value = value_arg_.GetValue(); + + auto compare_length = [op_type, value](int length) { + switch (op_type) { + case proto::plan::OpType::Equal: + return length == value; + case proto::plan::OpType::NotEqual: + return length != value; + case proto::plan::OpType::GreaterThan: + return length > value; + case proto::plan::OpType::GreaterEqual: + return length >= value; + case proto::plan::OpType::LessThan: + return length < value; + case proto::plan::OpType::LessEqual: + return length <= value; + default: + ThrowInfo(OpTypeInvalid, + "unsupported operator type for vector array " + "length eval expr: {}", + op_type); + } + return false; + }; + + auto execute_sub_batch = [compare_length]( + const VectorArrayView* data, + const bool* valid_data, + const int32_t* offsets, + const int size, + TargetBitmapView res, + TargetBitmapView valid_res) { + if (data == nullptr) { + return; + } + for (size_t i = 0; i < size; ++i) { + auto offset = i; + if constexpr (filter_type == FilterType::random) { + offset = (offsets) ? offsets[i] : i; + } + if (valid_data != nullptr && !valid_data[offset]) { + res[i] = valid_res[i] = false; + continue; + } + res[i] = compare_length(data[offset].length()); + } + }; + + int64_t processed_size = 0; + if (has_offset_input_) { + processed_size = ProcessDataByOffsets( + execute_sub_batch, std::nullptr_t{}, input, res, valid_res); + } else { + processed_size = ProcessDataChunks( + execute_sub_batch, std::nullptr_t{}, res, valid_res); + } + + AssertInfo(processed_size == real_batch_size, + "internal error: expr processed rows {} not equal " + "expect batch size {}", + processed_size, + real_batch_size); + return res_vec; +} + template VectorPtr PhyBinaryArithOpEvalRangeExpr::ExecRangeVisitorImpl(OffsetVector* input) { diff --git a/internal/core/src/exec/expression/BinaryArithOpEvalRangeExpr.h b/internal/core/src/exec/expression/BinaryArithOpEvalRangeExpr.h index a78c7b2d4a..edc2ad7ec3 100644 --- a/internal/core/src/exec/expression/BinaryArithOpEvalRangeExpr.h +++ b/internal/core/src/exec/expression/BinaryArithOpEvalRangeExpr.h @@ -518,8 +518,9 @@ class PhyBinaryArithOpEvalRangeExpr : public SegmentExpr { data_type = expr_->column_.element_type_; } - // JSON and ARRAY types cannot use index for arith ops - if (data_type == DataType::JSON || data_type == DataType::ARRAY) { + // JSON, ARRAY and VECTOR_ARRAY types cannot use index for arith ops. + if (data_type == DataType::JSON || data_type == DataType::ARRAY || + data_type == DataType::VECTOR_ARRAY) { exec_path_ = ExprExecPath::RawData; return; } @@ -592,6 +593,10 @@ class PhyBinaryArithOpEvalRangeExpr : public SegmentExpr { VectorPtr ExecRangeVisitorImplForArray(OffsetVector* input = nullptr); + template + VectorPtr + ExecRangeVisitorImplForVectorArray(OffsetVector* input = nullptr); + private: std::shared_ptr expr_; SingleElement right_operand_arg_; diff --git a/internal/core/src/exec/expression/Expr.h b/internal/core/src/exec/expression/Expr.h index 09e860b060..013339bffa 100644 --- a/internal/core/src/exec/expression/Expr.h +++ b/internal/core/src/exec/expression/Expr.h @@ -751,16 +751,50 @@ class SegmentExpr : public Expr { int64_t processed_size = 0; // index reverse lookup (only for ScalarIndex path) - if (UseIndexCursor() && num_data_chunk_ == 0) { - return ProcessIndexLookupByOffsets( - func, skip_func, input, res, valid_res, values...); + if constexpr (!std::is_same_v) { + if (UseIndexCursor() && num_data_chunk_ == 0) { + return ProcessIndexLookupByOffsets( + func, skip_func, input, res, valid_res, values...); + } } auto& skip_index = segment_->GetSkipIndex(); - // raw data scan - // sealed segment - if (segment_->type() == SegmentType::Sealed) { + if constexpr (std::is_same_v) { + for (size_t i = 0; i < input->size(); ++i) { + int64_t offset = (*input)[i]; + auto [chunk_id, chunk_offset] = + segment_->get_chunk_by_offset(field_id_, offset); + // chunk_data would read the wrong layout: + // storage holds VectorArray, and nullable rows may be compacted. + // Use chunk_view to build logical VectorArrayView rows. + auto pw = segment_->chunk_view( + op_ctx_, + field_id_, + chunk_id, + std::make_pair(chunk_offset, int64_t{1})); + const auto& [data_vec, valid_data] = pw.get(); + if (!skip_func || !skip_func(skip_index, field_id_, chunk_id)) { + func.template operator()( + data_vec.data(), + valid_data.data(), + nullptr, + 1, + res + processed_size, + valid_res + processed_size, + values...); + } else { + ApplyValidData(valid_data.data(), + res + processed_size, + valid_res + processed_size, + 1); + } + processed_size++; + } + return input->size(); + } else if (segment_->type() == SegmentType::Sealed) { + // raw data scan + // sealed segment if (segment_->is_chunked()) { if constexpr (std::is_same_v || std::is_same_v || @@ -1341,40 +1375,36 @@ class SegmentExpr : public Expr { continue; //do not go empty-loop at the bound of the chunk auto& skip_index = segment_->GetSkipIndex(); - auto pw = segment_->chunk_data(op_ctx_, field_id_, i); - auto chunk = pw.get(); - const bool* valid_data = chunk.valid_data(); - if (valid_data != nullptr) { - valid_data += data_pos; - } - if (!skip_func || !skip_func(skip_index, field_id_, i)) { - const T* data = chunk.data() + data_pos; - - if constexpr (NeedSegmentOffsets) { - // For GIS functions: construct segment offsets array - std::vector segment_offsets_array(size); - for (int64_t j = 0; j < size; ++j) { - segment_offsets_array[j] = static_cast( - size_per_chunk_ * i + data_pos + j); + auto process_chunk = [&](const T* data, const bool* valid_data) { + auto skipped = skip_func && skip_func(skip_index, field_id_, i); + if (!skipped) { + if constexpr (NeedSegmentOffsets) { + // For GIS functions: construct segment offsets array + std::vector segment_offsets_array(size); + for (int64_t j = 0; j < size; ++j) { + segment_offsets_array[j] = static_cast( + size_per_chunk_ * i + data_pos + j); + } + func(data, + valid_data, + nullptr, + segment_offsets_array.data(), + size, + res + processed_size, + valid_res + processed_size, + values...); + } else { + func(data, + valid_data, + nullptr, + size, + res + processed_size, + valid_res + processed_size, + values...); } - func(data, - valid_data, - nullptr, - segment_offsets_array.data(), - size, - res + processed_size, - valid_res + processed_size, - values...); - } else { - func(data, - valid_data, - nullptr, - size, - res + processed_size, - valid_res + processed_size, - values...); + return; } - } else { + // Chunk is skipped by SkipIndex. // We still need to: // 1. Apply valid_data to handle nullable fields @@ -1384,7 +1414,6 @@ class SegmentExpr : public Expr { res + processed_size, valid_res + processed_size, size); - // Call func with nullptr to update internal cursors if constexpr (NeedSegmentOffsets) { std::vector segment_offsets_array(size); for (int64_t j = 0; j < size; ++j) { @@ -1408,6 +1437,24 @@ class SegmentExpr : public Expr { valid_res + processed_size, values...); } + }; + + if constexpr (std::is_same_v) { + // chunk_data would read the wrong layout: + // storage holds VectorArray, and nullable rows may be compacted. + // Use chunk_view to build logical VectorArrayView rows. + auto pw = segment_->chunk_view( + op_ctx_, field_id_, i, std::make_pair(data_pos, size)); + const auto& [data_vec, valid_data] = pw.get(); + process_chunk(data_vec.data(), valid_data.data()); + } else { + auto pw = segment_->chunk_data(op_ctx_, field_id_, i); + auto chunk = pw.get(); + const bool* valid_data = chunk.valid_data(); + if (valid_data != nullptr) { + valid_data += data_pos; + } + process_chunk(chunk.data() + data_pos, valid_data); } processed_size += size; @@ -1467,7 +1514,8 @@ class SegmentExpr : public Expr { bool is_seal = false; if constexpr (std::is_same_v || std::is_same_v || - std::is_same_v) { + std::is_same_v || + std::is_same_v) { if (segment_->type() == SegmentType::Sealed) { // first is the raw data, second is valid_data // use valid_data to see if raw data is null @@ -1497,33 +1545,40 @@ class SegmentExpr : public Expr { is_seal = true; } } - if (!is_seal) { - auto pw = segment_->chunk_data(op_ctx_, field_id_, i); - auto chunk = pw.get(); - const T* data = chunk.data() + data_pos; - const bool* valid_data = chunk.valid_data(); - if (valid_data != nullptr) { - valid_data += data_pos; - } + if constexpr (std::is_same_v) { + AssertInfo(is_seal, + "VectorArrayView must be read through chunk " + "views"); + } else { + if (!is_seal) { + auto pw = + segment_->chunk_data(op_ctx_, field_id_, i); + auto chunk = pw.get(); + const T* data = chunk.data() + data_pos; + const bool* valid_data = chunk.valid_data(); + if (valid_data != nullptr) { + valid_data += data_pos; + } - if constexpr (NeedSegmentOffsets) { - // For GIS functions: construct segment offsets array - func(data, - valid_data, - nullptr, - segment_offsets_array.data(), - size, - res + processed_size, - valid_res + processed_size, - values...); - } else { - func(data, - valid_data, - nullptr, - size, - res + processed_size, - valid_res + processed_size, - values...); + if constexpr (NeedSegmentOffsets) { + // For GIS functions: construct segment offsets array + func(data, + valid_data, + nullptr, + segment_offsets_array.data(), + size, + res + processed_size, + valid_res + processed_size, + values...); + } else { + func(data, + valid_data, + nullptr, + size, + res + processed_size, + valid_res + processed_size, + values...); + } } } } else { @@ -1535,7 +1590,8 @@ class SegmentExpr : public Expr { const bool* valid_data; if constexpr (std::is_same_v || std::is_same_v || - std::is_same_v) { + std::is_same_v || + std::is_same_v) { auto pw = segment_->get_batch_views( op_ctx_, field_id_, i, data_pos, size); valid_data = pw.get().second.data(); diff --git a/internal/core/src/exec/expression/ExprArrayTest.cpp b/internal/core/src/exec/expression/ExprArrayTest.cpp index bce8ecfed2..6cc5037065 100644 --- a/internal/core/src/exec/expression/ExprArrayTest.cpp +++ b/internal/core/src/exec/expression/ExprArrayTest.cpp @@ -769,6 +769,142 @@ TEST(Expr, TestVectorArrayNullExpr) { (void)fakevec_fid; } +TEST(Expr, TestVectorArrayLengthExpr) { + auto schema = std::make_shared(); + auto fakevec_fid = schema->AddDebugField( + "fakevec", DataType::VECTOR_FLOAT, 16, knowhere::metric::L2); + auto i64_fid = schema->AddDebugField("id", DataType::INT64); + auto vector_array_fid = + schema->AddDebugVectorArrayField("array_float_vector", + DataType::VECTOR_FLOAT, + 4, + knowhere::metric::L2, + true); + schema->set_primary_field_id(i64_fid); + + constexpr int N = 128; + constexpr int kArrayLen = 2; + auto raw_data = DataGen(schema, N, 42, 0, 1, kArrayLen); + auto valid_data = raw_data.get_col_valid(vector_array_fid); + auto vector_array_col = + raw_data.get_col(vector_array_fid); + + auto growing = CreateGrowingSegment(schema, empty_index_meta); + auto offset = growing->PreInsert(N); + growing->Insert(offset, + N, + raw_data.row_ids_.data(), + raw_data.timestamps_.data(), + raw_data.raw_); + auto growing_segment = dynamic_cast(growing.get()); + ASSERT_NE(growing_segment, nullptr); + + std::vector valid_bitmap((N + 7) / 8, 0); + std::vector vector_arrays; + vector_arrays.reserve(N); + for (int i = 0; i < N; ++i) { + if (valid_data[i]) { + valid_bitmap[i >> 3] |= 1 << (i & 0x07); + vector_arrays.emplace_back(vector_array_col[i]); + } + } + + auto field_data = storage::CreateFieldData( + DataType::VECTOR_ARRAY, DataType::VECTOR_FLOAT, true, 4); + field_data->FillFieldData(vector_arrays.data(), valid_bitmap.data(), N, 0); + + auto cm = storage::RemoteChunkManagerSingleton::GetInstance() + .GetRemoteChunkManager(); + auto sealed_segment = CreateSealedSegment(schema); + auto field_data_info = + PrepareSingleFieldInsertBinlog(kCollectionID, + kPartitionID, + kSegmentID, + vector_array_fid.get(), + {field_data}, + cm); + sealed_segment->LoadFieldData(field_data_info); + + auto make_plan = [&](proto::plan::OpType op, int64_t target) { + proto::plan::GenericValue value; + value.set_int64_val(target); + proto::plan::GenericValue right_operand; + right_operand.set_int64_val(0); + auto expr = std::make_shared( + milvus::expr::ColumnInfo(vector_array_fid, + DataType::VECTOR_ARRAY, + DataType::VECTOR_FLOAT, + {}, + true), + op, + proto::plan::ArithOpType::ArrayLength, + value, + right_operand); + return std::make_shared(DEFAULT_PLANNODE_ID, + expr); + }; + + struct Testcase { + proto::plan::OpType op; + int64_t target; + std::function ref_func; + }; + std::vector testcases = { + {proto::plan::OpType::Equal, + kArrayLen, + [](bool valid) { return valid; }}, + {proto::plan::OpType::NotEqual, + kArrayLen + 1, + [](bool valid) { return valid; }}, + {proto::plan::OpType::GreaterThan, + kArrayLen - 1, + [](bool valid) { return valid; }}, + {proto::plan::OpType::GreaterEqual, + kArrayLen, + [](bool valid) { return valid; }}, + {proto::plan::OpType::LessThan, + kArrayLen, + [](bool valid) { return false; }}, + {proto::plan::OpType::LessEqual, + kArrayLen, + [](bool valid) { return valid; }}, + }; + + for (const auto& testcase : testcases) { + std::array segments = { + static_cast(growing_segment), + static_cast(sealed_segment.get())}; + for (auto* segment : segments) { + auto plan = make_plan(testcase.op, testcase.target); + auto final = ExecuteQueryExpr(plan, segment, N, MAX_TIMESTAMP); + EXPECT_EQ(final.size(), N); + for (int i = 0; i < N; ++i) { + ASSERT_EQ(final[i], testcase.ref_func(valid_data[i])) + << "segment type " << segment->type() << ", row " << i; + } + + milvus::exec::OffsetVector offsets; + offsets.reserve(N / 2); + for (int i = 0; i < N; ++i) { + if (i % 2 == 0) { + offsets.emplace_back(i); + } + } + auto col_vec = milvus::test::gen_filter_res( + plan.get(), segment, N, MAX_TIMESTAMP, &offsets); + BitsetTypeView view(col_vec->GetRawData(), col_vec->size()); + ASSERT_EQ(view.size(), offsets.size()); + for (int i = 0; i < offsets.size(); ++i) { + ASSERT_EQ(view[i], testcase.ref_func(valid_data[offsets[i]])) + << "segment type " << segment->type() << ", offset row " + << offsets[i]; + } + } + } + + (void)fakevec_fid; +} + TEST(Expr, PraseArrayContainsExpr) { // Test expressions for array_contains operations std::vector exprs{ diff --git a/internal/core/src/segcore/SegmentGrowingImpl.cpp b/internal/core/src/segcore/SegmentGrowingImpl.cpp index c9b5e8839d..41af60dbdb 100644 --- a/internal/core/src/segcore/SegmentGrowingImpl.cpp +++ b/internal/core/src/segcore/SegmentGrowingImpl.cpp @@ -1284,8 +1284,90 @@ SegmentGrowingImpl::chunk_vector_array_view_impl( FieldId field_id, int64_t chunk_id, std::optional> offset_len) const { - ThrowInfo(ErrorCode::NotImplemented, - "chunk vector array view impl not implement for growing segment"); + (void)op_ctx; + + auto& field_meta = schema_->operator[](field_id); + AssertInfo(field_meta.get_data_type() == DataType::VECTOR_ARRAY, + "chunk_vector_array_view_impl only supports VECTOR_ARRAY field"); + + auto vector_data = insert_record_.get_data(field_id); + auto size_per_chunk = vector_data->get_size_per_chunk(); + auto active_count = insert_record_.ack_responder_.GetAck(); + auto start_offset = int64_t{0}; + auto len = size_per_chunk; + if (offset_len.has_value()) { + start_offset = offset_len->first; + len = offset_len->second; + } + + AssertInfo(start_offset >= 0 && start_offset < size_per_chunk, + "Retrieve vector array views with out-of-bound offset:{}, " + "len:{}, wrong", + start_offset, + len); + AssertInfo(len >= 0 && len <= size_per_chunk, + "Retrieve vector array views with out-of-bound offset:{}, " + "len:{}, wrong", + start_offset, + len); + AssertInfo(start_offset + len <= size_per_chunk, + "Retrieve vector array views with out-of-bound offset:{}, " + "len:{}, wrong", + start_offset, + len); + + auto logical_start = chunk_id * size_per_chunk + start_offset; + AssertInfo(logical_start >= 0 && logical_start <= active_count, + "Retrieve vector array views with out-of-bound offset:{}, " + "len:{}, wrong", + start_offset, + len); + if (offset_len.has_value()) { + AssertInfo(logical_start + len <= active_count, + "Retrieve vector array views with out-of-bound offset:{}, " + "len:{}, wrong", + start_offset, + len); + } else { + len = std::min(len, active_count - logical_start); + } + + std::vector views; + views.reserve(len); + FixedVector valid_data; + ThreadSafeValidDataPtr valid_vec_ptr = nullptr; + if (field_meta.is_nullable()) { + valid_vec_ptr = insert_record_.get_valid_data(field_id); + valid_data.reserve(len); + } + + for (int64_t i = 0; i < len; ++i) { + auto logical_offset = logical_start + i; + if (field_meta.is_nullable()) { + auto valid = valid_vec_ptr->is_valid(logical_offset); + valid_data.push_back(valid); + if (!valid) { + views.emplace_back(); + continue; + } + } + + auto vector_array = vector_data->get_element(logical_offset); + AssertInfo(vector_array != nullptr, + "Cannot find VECTOR_ARRAY data at segment offset {}", + logical_offset); + views.emplace_back(const_cast(vector_array->data()), + vector_array->dim(), + vector_array->length(), + vector_array->byte_size(), + vector_array->get_element_type()); + } + + std::pair, FixedVector> content{ + std::move(views), std::move(valid_data)}; + return PinWrapper< + std::pair, FixedVector>>( + std::move(content)); } PinWrapper, FixedVector>> diff --git a/internal/parser/planparserv2/parser_visitor.go b/internal/parser/planparserv2/parser_visitor.go index f7f89b7b45..2567963d9f 100644 --- a/internal/parser/planparserv2/parser_visitor.go +++ b/internal/parser/planparserv2/parser_visitor.go @@ -2281,17 +2281,28 @@ func (v *ParserVisitor) VisitArrayLength(ctx *parser.ArrayLengthContext) interfa FieldId: field.FieldID, DataType: field.DataType, ElementType: field.GetElementType(), + Nullable: field.GetNullable(), } } else { - columnInfo, err = v.getChildColumnInfo(ctx.Identifier(), ctx.JSONIdentifier(), nil, nil) + if ctx.Identifier() != nil { + if parentColumnInfo, ok, parentErr := v.getStructArrayParentColumnInfo(ctx.Identifier().GetText()); ok || parentErr != nil { + columnInfo = parentColumnInfo + err = parentErr + } + } + if columnInfo == nil && err == nil { + columnInfo, err = v.getChildColumnInfo(ctx.Identifier(), ctx.JSONIdentifier(), nil, nil) + } if err != nil { return err } } if columnInfo == nil || - (!typeutil.IsJSONType(columnInfo.GetDataType()) && !typeutil.IsArrayType(columnInfo.GetDataType())) { + (!typeutil.IsJSONType(columnInfo.GetDataType()) && + !typeutil.IsArrayType(columnInfo.GetDataType()) && + !typeutil.IsVectorArrayType(columnInfo.GetDataType())) { return merr.WrapErrParameterInvalidMsg( - "array_length operation are only supported on json or array fields now, got: %s", ctx.GetText()) + "array_length operation are only supported on json, array or array-of-vector fields now, got: %s", ctx.GetText()) } expr := &planpb.Expr{ diff --git a/internal/parser/planparserv2/plan_parser_v2_test.go b/internal/parser/planparserv2/plan_parser_v2_test.go index 7899b3c0ee..3e8669348d 100644 --- a/internal/parser/planparserv2/plan_parser_v2_test.go +++ b/internal/parser/planparserv2/plan_parser_v2_test.go @@ -990,6 +990,7 @@ func TestExpr_VectorArrayOnlyStructParentIsNull(t *testing.T) { Name: "vector_struct[embeddings]", DataType: schemapb.DataType_ArrayOfVector, ElementType: schemapb.DataType_FloatVector, + Nullable: true, TypeParams: []*commonpb.KeyValuePair{ {Key: common.DimKey, Value: "4"}, }, @@ -2149,6 +2150,7 @@ func Test_ArrayLength(t *testing.T) { `array_length(StringArrayField) > 5`, `array_length(StringArrayField) >= 5`, // struct array sub-field + `array_length(struct_array) == 2`, `array_length(struct_array[sub_str]) == 3`, `array_length(struct_array[sub_int]) > 1`, `array_length(struct_array[sub_int]) <= 10`, @@ -3188,7 +3190,7 @@ func TestExpr_StructIndexField_PlanShape(t *testing.T) { func TestExpr_StructFieldArrayLength(t *testing.T) { schema := newTestSchemaHelper(t) - ret := handleExpr(schema, `array_length(struct_array[sub_int])`) + ret := handleExpr(schema, `array_length(struct_array)`) ewt, ok := ret.(*ExprWithType) require.True(t, ok, "handleExpr should return *ExprWithType") assert.Equal(t, schemapb.DataType_Int64, ewt.dataType) @@ -3200,12 +3202,30 @@ func TestExpr_StructFieldArrayLength(t *testing.T) { columnInfo := bae.GetLeft().GetColumnExpr().GetInfo() require.NotNil(t, columnInfo) + assert.Equal(t, int64(133), columnInfo.GetFieldId()) + assert.Equal(t, schemapb.DataType_Array, columnInfo.GetDataType()) + assert.Equal(t, schemapb.DataType_VarChar, columnInfo.GetElementType()) + assert.Empty(t, columnInfo.GetNestedPath()) + + ret = handleExpr(schema, `array_length(struct_array[sub_int])`) + ewt, ok = ret.(*ExprWithType) + require.True(t, ok, "handleExpr should return *ExprWithType") + assert.Equal(t, schemapb.DataType_Int64, ewt.dataType) + + bae = ewt.expr.GetBinaryArithExpr() + require.NotNil(t, bae) + assert.Equal(t, planpb.ArithOpType_ArrayLength, bae.GetOp()) + require.Nil(t, bae.GetRight()) + + columnInfo = bae.GetLeft().GetColumnExpr().GetInfo() + require.NotNil(t, columnInfo) assert.Equal(t, int64(134), columnInfo.GetFieldId()) assert.Equal(t, schemapb.DataType_Array, columnInfo.GetDataType()) assert.Equal(t, schemapb.DataType_Int32, columnInfo.GetElementType()) assert.Empty(t, columnInfo.GetNestedPath()) validExprs := []string{ + `array_length(struct_array) == 10`, `array_length(struct_array[sub_int]) == 10`, `array_length(struct_array[sub_str]) > 0`, } @@ -3220,6 +3240,58 @@ func TestExpr_StructFieldArrayLength(t *testing.T) { } } +func TestExpr_VectorArrayOnlyStructParentArrayLength(t *testing.T) { + schema := newTestSchema(false) + schema.StructArrayFields = append(schema.StructArrayFields, &schemapb.StructArrayFieldSchema{ + FieldID: 10000, + Name: "vector_struct", + Nullable: true, + Fields: []*schemapb.FieldSchema{ + { + FieldID: 10001, + Name: "vector_struct[embeddings]", + DataType: schemapb.DataType_ArrayOfVector, + ElementType: schemapb.DataType_FloatVector, + Nullable: true, + TypeParams: []*commonpb.KeyValuePair{ + {Key: common.DimKey, Value: "4"}, + }, + }, + }, + }) + helper, err := typeutil.CreateSchemaHelper(schema) + require.NoError(t, err) + + for _, expr := range []string{ + `array_length(vector_struct)`, + `array_length(vector_struct[embeddings])`, + } { + ret := handleExpr(helper, expr) + ewt, ok := ret.(*ExprWithType) + require.True(t, ok, "handleExpr should return *ExprWithType") + assert.Equal(t, schemapb.DataType_Int64, ewt.dataType) + + bae := ewt.expr.GetBinaryArithExpr() + require.NotNil(t, bae) + assert.Equal(t, planpb.ArithOpType_ArrayLength, bae.GetOp()) + + columnInfo := bae.GetLeft().GetColumnExpr().GetInfo() + require.NotNil(t, columnInfo) + assert.Equal(t, int64(10001), columnInfo.GetFieldId()) + assert.Equal(t, schemapb.DataType_ArrayOfVector, columnInfo.GetDataType()) + assert.Equal(t, schemapb.DataType_FloatVector, columnInfo.GetElementType()) + assert.True(t, columnInfo.GetNullable()) + } + + _, err = CreateSearchPlan(helper, `array_length(vector_struct) == 2`, "FloatVectorField", &planpb.QueryInfo{ + Topk: 0, + MetricType: "", + SearchParams: "", + RoundDecimal: 0, + }, nil, nil) + require.NoError(t, err) +} + func TestExpr_StructIndexField_RangeForms(t *testing.T) { schema := newTestSchemaHelper(t)