furyhawk/milvus
1
0
Fork 0
mirror of https://github.com/milvus-io/milvus.git synced 2026-07-21 10:15:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add timesatmptz data type (#44005)

Browse Source 
issue: https://github.com/milvus-io/milvus/issues/27467
>
https://github.com/milvus-io/milvus/issues/27467#issuecomment-3092211420
> * [x]  M1 Create collection with timestamptz field
> * [x]  M2 Insert timestamptz field data
> * [x]  M3 Retrieve timestamptz field data
> * [x]  M4 Implement handoff[ ]  

The second PR of issue:
https://github.com/milvus-io/milvus/issues/27467, which completes M1-M4
described above.

---------

Signed-off-by: xtx <xtianx@smail.nju.edu.cn>
This commit is contained in:
Tianx
2025-08-26 15:59:53 +08:00
committed by GitHub
parent f301692900
commit c0d62268ac
72 changed files with 1371 additions and 354 deletions
Download Patch File Download Diff File Expand all files Collapse all files
client/column/columns.go
+3
View File
@@ -196,6 +196,9 @@ func FieldDataColumn(fd *schemapb.FieldData, begin, end int) (Column, error) {
case schemapb.DataType_Double:
return parseScalarData(fd.GetFieldName(), fd.GetScalars().GetDoubleData().GetData(), begin, end, validData, NewColumnDouble, NewNullableColumnDouble)
case schemapb.DataType_Timestamptz:
return parseScalarData(fd.GetFieldName(), fd.GetScalars().GetTimestamptzData().GetData(), begin, end, validData, NewColumnTimestamptz, NewNullableColumnTimestamptz)
case schemapb.DataType_String:
return parseScalarData(fd.GetFieldName(), fd.GetScalars().GetStringData().GetData(), begin, end, validData, NewColumnString, NewNullableColumnString)
client/column/nullable.go
+11 -10
View File
@@ -18,16 +18,17 @@ package column
var (
// scalars
NewNullableColumnBool NullableColumnCreateFunc[bool, *ColumnBool] = NewNullableColumnCreator(NewColumnBool).New
NewNullableColumnInt8 NullableColumnCreateFunc[int8, *ColumnInt8] = NewNullableColumnCreator(NewColumnInt8).New
NewNullableColumnInt16 NullableColumnCreateFunc[int16, *ColumnInt16] = NewNullableColumnCreator(NewColumnInt16).New
NewNullableColumnInt32 NullableColumnCreateFunc[int32, *ColumnInt32] = NewNullableColumnCreator(NewColumnInt32).New
NewNullableColumnInt64 NullableColumnCreateFunc[int64, *ColumnInt64] = NewNullableColumnCreator(NewColumnInt64).New
NewNullableColumnVarChar NullableColumnCreateFunc[string, *ColumnVarChar] = NewNullableColumnCreator(NewColumnVarChar).New
NewNullableColumnString NullableColumnCreateFunc[string, *ColumnString] = NewNullableColumnCreator(NewColumnString).New
NewNullableColumnFloat NullableColumnCreateFunc[float32, *ColumnFloat] = NewNullableColumnCreator(NewColumnFloat).New
NewNullableColumnDouble NullableColumnCreateFunc[float64, *ColumnDouble] = NewNullableColumnCreator(NewColumnDouble).New
NewNullableColumnJSONBytes NullableColumnCreateFunc[[]byte, *ColumnJSONBytes] = NewNullableColumnCreator(NewColumnJSONBytes).New
NewNullableColumnBool NullableColumnCreateFunc[bool, *ColumnBool] = NewNullableColumnCreator(NewColumnBool).New
NewNullableColumnInt8 NullableColumnCreateFunc[int8, *ColumnInt8] = NewNullableColumnCreator(NewColumnInt8).New
NewNullableColumnInt16 NullableColumnCreateFunc[int16, *ColumnInt16] = NewNullableColumnCreator(NewColumnInt16).New
NewNullableColumnInt32 NullableColumnCreateFunc[int32, *ColumnInt32] = NewNullableColumnCreator(NewColumnInt32).New
NewNullableColumnInt64 NullableColumnCreateFunc[int64, *ColumnInt64] = NewNullableColumnCreator(NewColumnInt64).New
NewNullableColumnVarChar NullableColumnCreateFunc[string, *ColumnVarChar] = NewNullableColumnCreator(NewColumnVarChar).New
NewNullableColumnString NullableColumnCreateFunc[string, *ColumnString] = NewNullableColumnCreator(NewColumnString).New
NewNullableColumnFloat NullableColumnCreateFunc[float32, *ColumnFloat] = NewNullableColumnCreator(NewColumnFloat).New
NewNullableColumnDouble NullableColumnCreateFunc[float64, *ColumnDouble] = NewNullableColumnCreator(NewColumnDouble).New
NewNullableColumnTimestamptz NullableColumnCreateFunc[int64, *ColumnTimestamptz] = NewNullableColumnCreator(NewColumnTimestamptz).New
NewNullableColumnJSONBytes NullableColumnCreateFunc[[]byte, *ColumnJSONBytes] = NewNullableColumnCreator(NewColumnJSONBytes).New
// array
NewNullableColumnBoolArray NullableColumnCreateFunc[[]bool, *ColumnBoolArray] = NewNullableColumnCreator(NewColumnBoolArray).New
NewNullableColumnInt8Array NullableColumnCreateFunc[[]int8, *ColumnInt8Array] = NewNullableColumnCreator(NewColumnInt8Array).New
client/column/scalar.go
+17
View File
@@ -208,6 +208,23 @@ func (c *ColumnDouble) Slice(start, end int) Column {
}
}
/* Timestamptz */
var _ Column = (*ColumnTimestamptz)(nil)
type ColumnTimestamptz struct {
*genericColumnBase[int64]
}
func NewColumnTimestamptz(name string, values []int64) *ColumnTimestamptz {
return &ColumnTimestamptz{
genericColumnBase: &genericColumnBase[int64]{
name: name,
fieldType: entity.FieldTypeTimestamptz,
values: values,
},
}
}
/* Varchar */
var _ (Column) = (*ColumnVarChar)(nil)
client/entity/field.go
+17
View File
@@ -46,6 +46,8 @@ func (t FieldType) Name() string {
return "Float"
case FieldTypeDouble:
return "Double"
case FieldTypeTimestamptz:
return "Timestamptz"
case FieldTypeString:
return "String"
case FieldTypeVarChar:
@@ -86,6 +88,8 @@ func (t FieldType) String() string {
return "float32"
case FieldTypeDouble:
return "float64"
case FieldTypeTimestamptz:
return "timestamptz"
case FieldTypeString:
return "string"
case FieldTypeVarChar:
@@ -126,6 +130,8 @@ func (t FieldType) PbFieldType() (string, string) {
return "Float", "float32"
case FieldTypeDouble:
return "Double", "float64"
case FieldTypeTimestamptz:
return "Timestamptz", "int64" // Timestamptz
case FieldTypeString:
return "String", "string"
case FieldTypeVarChar:
@@ -165,6 +171,8 @@ const (
FieldTypeFloat FieldType = 10
// FieldTypeDouble field type double
FieldTypeDouble FieldType = 11
// FieldTypeTimestamptz field type timestamptz
FieldTypeTimestamptz FieldType = 15
// FieldTypeString field type string
FieldTypeString FieldType = 20
// FieldTypeVarChar field type varchar
@@ -323,6 +331,15 @@ func (f *Field) WithDefaultValueDouble(defaultValue float64) *Field {
return f
}
func (f *Field) WithDefaultValueTimestamptz(defaultValue int64) *Field {
f.DefaultValue = &schemapb.ValueField{
Data: &schemapb.ValueField_TimestamptzData{
TimestamptzData: defaultValue,
},
}
return f
}
func (f *Field) WithDefaultValueString(defaultValue string) *Field {
f.DefaultValue = &schemapb.ValueField{
Data: &schemapb.ValueField_StringData{
internal/core/src/common/ChunkWriter.cpp
+3
View File
@@ -414,6 +414,9 @@ create_chunk_writer(const FieldMeta& field_meta, Args&&... args) {
case milvus::DataType::DOUBLE:
return std::make_shared<ChunkWriter<arrow::DoubleArray, double>>(
dim, std::forward<Args>(args)..., nullable);
case milvus::DataType::TIMESTAMPTZ:
return std::make_shared<ChunkWriter<arrow::Int64Array, int64_t>>(
dim, std::forward<Args>(args)..., nullable);
case milvus::DataType::VECTOR_FLOAT:
return std::make_shared<
ChunkWriter<arrow::FixedSizeBinaryArray, knowhere::fp32>>(
internal/core/src/common/FieldData.cpp
+26
View File
@@ -209,6 +209,18 @@ FieldDataImpl<Type, is_type_entire_row>::FillFieldData(
}
return FillFieldData(array_info.first, array_info.second);
}
case DataType::TIMESTAMPTZ: {
auto array_info =
GetDataInfoFromArray<arrow::Int64Array,
arrow::Type::type::INT64>(array);
if (nullable_) {
return FillFieldData(array_info.first,
array->null_bitmap_data(),
element_count,
array->offset());
}
return FillFieldData(array_info.first, array_info.second);
}
case DataType::STRING:
case DataType::VARCHAR:
case DataType::TEXT: {
@@ -409,6 +421,17 @@ FieldDataImpl<Type, is_type_entire_row>::FillFieldData(
return FillFieldData(
values.data(), valid_data_ptr.get(), element_count, 0);
}
case DataType::TIMESTAMPTZ: {
FixedVector<int64_t> values(element_count);
if (default_value.has_value()) {
std::fill(values.begin(),
values.end(),
default_value->timestamptz_data());
return FillFieldData(values.data(), nullptr, element_count, 0);
}
return FillFieldData(
values.data(), valid_data_ptr.get(), element_count, 0);
}
case DataType::STRING:
case DataType::VARCHAR: {
FixedVector<std::string> values(element_count);
@@ -485,6 +508,9 @@ InitScalarFieldData(const DataType& type, bool nullable, int64_t cap_rows) {
case DataType::DOUBLE:
return std::make_shared<FieldData<double>>(
type, nullable, cap_rows);
case DataType::TIMESTAMPTZ:
return std::make_shared<FieldData<int64_t>>(
type, nullable, cap_rows);
case DataType::STRING:
case DataType::VARCHAR:
case DataType::TEXT:
internal/core/src/common/FieldMeta.cpp
+3
View File
@@ -11,12 +11,14 @@
#include "common/FieldMeta.h"
#include "common/SystemProperty.h"
#include "common/Types.h"
#include "common/protobuf_utils.h"
#include "common/Common.h"
#include <boost/lexical_cast.hpp>
#include <optional>
#include "Consts.h"
#include "log/Log.h"
namespace milvus {
TokenizerParams
@@ -90,6 +92,7 @@ FieldMeta::ParseFrom(const milvus::proto::schema::FieldSchema& schema_proto) {
return schema_proto.default_value();
}();
LOG_DEBUG("Get datatype {}", GetDataTypeName(data_type));
if (data_type == DataType::VECTOR_ARRAY) {
// todo(SpadeA): revisit the code when index build for vector array is ready
int64_t dim = 0;
internal/core/src/common/Types.h
+19
View File
@@ -83,6 +83,7 @@ enum class DataType {
JSON = 23,
// GEOMETRY = 24 // reserved in proto
TEXT = 25,
TIMESTAMPTZ = 26, // Timestamp with timezone, stored as int64
// Some special Data type, start from after 50
// just for internal use now, may sync proto in future
@@ -128,6 +129,8 @@ GetDataTypeSize(DataType data_type, int dim = 1) {
return sizeof(float);
case DataType::DOUBLE:
return sizeof(double);
case DataType::TIMESTAMPTZ:
return sizeof(int64_t);
case DataType::VECTOR_FLOAT:
return sizeof(float) * dim;
case DataType::VECTOR_BINARY: {
@@ -170,6 +173,8 @@ GetArrowDataType(DataType data_type, int dim = 1) {
return arrow::float32();
case DataType::DOUBLE:
return arrow::float64();
case DataType::TIMESTAMPTZ:
return arrow::int64();
case DataType::STRING:
case DataType::VARCHAR:
case DataType::TEXT:
@@ -227,6 +232,8 @@ GetDataTypeName(DataType data_type) {
return "float";
case DataType::DOUBLE:
return "double";
case DataType::TIMESTAMPTZ:
return "timestamptz";
case DataType::STRING:
return "string";
case DataType::VARCHAR:
@@ -597,6 +604,15 @@ struct TypeTraits<DataType::DOUBLE> {
static constexpr const char* Name = "DOUBLE";
};
template <>
struct TypeTraits<DataType::TIMESTAMPTZ> {
using NativeType = double;
static constexpr DataType TypeKind = DataType::TIMESTAMPTZ;
static constexpr bool IsPrimitiveType = true;
static constexpr bool IsFixedWidth = true;
static constexpr const char* Name = "TIMESTAMPTZ";
};
template <>
struct TypeTraits<DataType::VARCHAR> {
using NativeType = std::string;
@@ -720,6 +736,9 @@ struct fmt::formatter<milvus::DataType> : formatter<string_view> {
case milvus::DataType::DOUBLE:
name = "DOUBLE";
break;
case milvus::DataType::TIMESTAMPTZ:
name = "TIMESTAMPTZ";
break;
case milvus::DataType::STRING:
name = "STRING";
break;
internal/core/src/exec/expression/UnaryExpr.cpp
+200 -205
View File
@@ -330,9 +330,8 @@ PhyUnaryRangeFilterExpr::ExecRangeVisitorImplArray(EvalCtx& context) {
}
int processed_cursor = 0;
auto execute_sub_batch =
[op_type,
&processed_cursor,
&bitmap_input]<FilterType filter_type = FilterType::sequential>(
[ op_type, &processed_cursor, &
bitmap_input ]<FilterType filter_type = FilterType::sequential>(
const milvus::ArrayView* data,
const bool* valid_data,
const int32_t* offsets,
@@ -341,186 +340,185 @@ PhyUnaryRangeFilterExpr::ExecRangeVisitorImplArray(EvalCtx& context) {
TargetBitmapView valid_res,
ValueType val,
int index) {
switch (op_type) {
case proto::plan::GreaterThan: {
UnaryElementFuncForArray<ValueType,
proto::plan::GreaterThan,
filter_type>
func;
func(data,
valid_data,
size,
val,
index,
res,
valid_res,
bitmap_input,
processed_cursor,
offsets);
break;
}
case proto::plan::GreaterEqual: {
UnaryElementFuncForArray<ValueType,
proto::plan::GreaterEqual,
filter_type>
func;
func(data,
valid_data,
size,
val,
index,
res,
valid_res,
bitmap_input,
processed_cursor,
offsets);
break;
}
case proto::plan::LessThan: {
UnaryElementFuncForArray<ValueType,
proto::plan::LessThan,
filter_type>
func;
func(data,
valid_data,
size,
val,
index,
res,
valid_res,
bitmap_input,
processed_cursor,
offsets);
break;
}
case proto::plan::LessEqual: {
UnaryElementFuncForArray<ValueType,
proto::plan::LessEqual,
filter_type>
func;
func(data,
valid_data,
size,
val,
index,
res,
valid_res,
bitmap_input,
processed_cursor,
offsets);
break;
}
case proto::plan::Equal: {
UnaryElementFuncForArray<ValueType,
proto::plan::Equal,
filter_type>
func;
func(data,
valid_data,
size,
val,
index,
res,
valid_res,
bitmap_input,
processed_cursor,
offsets);
break;
}
case proto::plan::NotEqual: {
UnaryElementFuncForArray<ValueType,
proto::plan::NotEqual,
filter_type>
func;
func(data,
valid_data,
size,
val,
index,
res,
valid_res,
bitmap_input,
processed_cursor,
offsets);
break;
}
case proto::plan::PrefixMatch: {
UnaryElementFuncForArray<ValueType,
proto::plan::PrefixMatch,
filter_type>
func;
func(data,
valid_data,
size,
val,
index,
res,
valid_res,
bitmap_input,
processed_cursor,
offsets);
break;
}
case proto::plan::Match: {
UnaryElementFuncForArray<ValueType,
proto::plan::Match,
filter_type>
func;
func(data,
valid_data,
size,
val,
index,
res,
valid_res,
bitmap_input,
processed_cursor,
offsets);
break;
}
case proto::plan::PostfixMatch: {
UnaryElementFuncForArray<ValueType,
proto::plan::PostfixMatch,
filter_type>
func;
func(data,
valid_data,
size,
val,
index,
res,
valid_res,
bitmap_input,
processed_cursor,
offsets);
break;
}
case proto::plan::InnerMatch: {
UnaryElementFuncForArray<ValueType,
proto::plan::InnerMatch,
filter_type>
func;
func(data,
valid_data,
size,
val,
index,
res,
valid_res,
bitmap_input,
processed_cursor,
offsets);
break;
}
default:
ThrowInfo(
OpTypeInvalid,
fmt::format(
"unsupported operator type for unary expr: {}",
op_type));
switch (op_type) {
case proto::plan::GreaterThan: {
UnaryElementFuncForArray<ValueType,
proto::plan::GreaterThan,
filter_type>
func;
func(data,
valid_data,
size,
val,
index,
res,
valid_res,
bitmap_input,
processed_cursor,
offsets);
break;
}
processed_cursor += size;
};
case proto::plan::GreaterEqual: {
UnaryElementFuncForArray<ValueType,
proto::plan::GreaterEqual,
filter_type>
func;
func(data,
valid_data,
size,
val,
index,
res,
valid_res,
bitmap_input,
processed_cursor,
offsets);
break;
}
case proto::plan::LessThan: {
UnaryElementFuncForArray<ValueType,
proto::plan::LessThan,
filter_type>
func;
func(data,
valid_data,
size,
val,
index,
res,
valid_res,
bitmap_input,
processed_cursor,
offsets);
break;
}
case proto::plan::LessEqual: {
UnaryElementFuncForArray<ValueType,
proto::plan::LessEqual,
filter_type>
func;
func(data,
valid_data,
size,
val,
index,
res,
valid_res,
bitmap_input,
processed_cursor,
offsets);
break;
}
case proto::plan::Equal: {
UnaryElementFuncForArray<ValueType,
proto::plan::Equal,
filter_type>
func;
func(data,
valid_data,
size,
val,
index,
res,
valid_res,
bitmap_input,
processed_cursor,
offsets);
break;
}
case proto::plan::NotEqual: {
UnaryElementFuncForArray<ValueType,
proto::plan::NotEqual,
filter_type>
func;
func(data,
valid_data,
size,
val,
index,
res,
valid_res,
bitmap_input,
processed_cursor,
offsets);
break;
}
case proto::plan::PrefixMatch: {
UnaryElementFuncForArray<ValueType,
proto::plan::PrefixMatch,
filter_type>
func;
func(data,
valid_data,
size,
val,
index,
res,
valid_res,
bitmap_input,
processed_cursor,
offsets);
break;
}
case proto::plan::Match: {
UnaryElementFuncForArray<ValueType,
proto::plan::Match,
filter_type>
func;
func(data,
valid_data,
size,
val,
index,
res,
valid_res,
bitmap_input,
processed_cursor,
offsets);
break;
}
case proto::plan::PostfixMatch: {
UnaryElementFuncForArray<ValueType,
proto::plan::PostfixMatch,
filter_type>
func;
func(data,
valid_data,
size,
val,
index,
res,
valid_res,
bitmap_input,
processed_cursor,
offsets);
break;
}
case proto::plan::InnerMatch: {
UnaryElementFuncForArray<ValueType,
proto::plan::InnerMatch,
filter_type>
func;
func(data,
valid_data,
size,
val,
index,
res,
valid_res,
bitmap_input,
processed_cursor,
offsets);
break;
}
default:
ThrowInfo(
OpTypeInvalid,
fmt::format("unsupported operator type for unary expr: {}",
op_type));
}
processed_cursor += size;
};
int64_t processed_size;
if (has_offset_input_) {
processed_size =
@@ -720,18 +718,16 @@ PhyUnaryRangeFilterExpr::ExecRangeVisitorImplJson(EvalCtx& context) {
} while (false)
int processed_cursor = 0;
auto execute_sub_batch = [op_type,
pointer,
&processed_cursor,
&bitmap_input]<FilterType filter_type =
FilterType::sequential>(
const milvus::Json* data,
const bool* valid_data,
const int32_t* offsets,
const int size,
TargetBitmapView res,
TargetBitmapView valid_res,
ExprValueType val) {
auto execute_sub_batch =
[ op_type, pointer, &processed_cursor, &
bitmap_input ]<FilterType filter_type = FilterType::sequential>(
const milvus::Json* data,
const bool* valid_data,
const int32_t* offsets,
const int size,
TargetBitmapView res,
TargetBitmapView valid_res,
ExprValueType val) {
bool has_bitmap_input = !bitmap_input.empty();
switch (op_type) {
case proto::plan::GreaterThan: {
@@ -1770,17 +1766,16 @@ PhyUnaryRangeFilterExpr::ExecRangeVisitorImplForData(EvalCtx& context) {
auto expr_type = expr_->op_type_;
size_t processed_cursor = 0;
auto execute_sub_batch = [expr_type,
&processed_cursor,
&bitmap_input]<FilterType filter_type =
FilterType::sequential>(
const T* data,
const bool* valid_data,
const int32_t* offsets,
const int size,
TargetBitmapView res,
TargetBitmapView valid_res,
IndexInnerType val) {
auto execute_sub_batch =
[ expr_type, &processed_cursor, &
bitmap_input ]<FilterType filter_type = FilterType::sequential>(
const T* data,
const bool* valid_data,
const int32_t* offsets,
const int size,
TargetBitmapView res,
TargetBitmapView valid_res,
IndexInnerType val) {
switch (expr_type) {
case proto::plan::GreaterThan: {
UnaryElementFunc<T, proto::plan::GreaterThan, filter_type> func;
internal/core/src/exec/operator/RescoresNode.cpp
+7 -6
View File
@@ -66,8 +66,9 @@ PhyRescoresNode::GetOutput() {
for (size_t i = 0; i < search_result.seg_offsets_.size(); i++) {
// remain offset will be -1 if result count not enough (less than topk)
// skip placeholder offset
if (search_result.seg_offsets_[i] >= 0){
offsets.push_back(static_cast<int32_t>(search_result.seg_offsets_[i]));
if (search_result.seg_offsets_[i] >= 0) {
offsets.push_back(
static_cast<int32_t>(search_result.seg_offsets_[i]));
offset_idx.push_back(i);
}
}
@@ -99,8 +100,8 @@ PhyRescoresNode::GetOutput() {
Assert(bitsetview.size() == offsets.size());
for (auto i = 0; i < offsets.size(); i++) {
if (bitsetview[i] > 0) {
search_result.distances_[offset_idx[i]] =
scorer->rescore(search_result.distances_[offset_idx[i]]);
search_result.distances_[offset_idx[i]] = scorer->rescore(
search_result.distances_[offset_idx[i]]);
}
}
} else {
@@ -115,8 +116,8 @@ PhyRescoresNode::GetOutput() {
bitset.append(view);
for (auto i = 0; i < offsets.size(); i++) {
if (bitset[offsets[i]] > 0) {
search_result.distances_[offset_idx[i]] =
scorer->rescore(search_result.distances_[offset_idx[i]]);
search_result.distances_[offset_idx[i]] = scorer->rescore(
search_result.distances_[offset_idx[i]]);
}
}
}
internal/core/src/exec/operator/Utils.h
+6 -2
View File
@@ -118,11 +118,15 @@ sort_search_result(milvus::SearchResult& result, bool large_is_better) {
if (large_is_better) {
std::sort(idx.begin(), idx.end(), [&](size_t i, size_t j) {
return result.distances_[i] > result.distances_[j] || (result.seg_offsets_[j] >=0 &&result.seg_offsets_[j] < 0);
return result.distances_[i] > result.distances_[j] ||
(result.seg_offsets_[j] >= 0 &&
result.seg_offsets_[j] < 0);
});
} else {
std::sort(idx.begin(), idx.end(), [&](size_t i, size_t j) {
return result.distances_[i] < result.distances_[j] || (result.seg_offsets_[j] >=0 &&result.seg_offsets_[j] < 0);
return result.distances_[i] < result.distances_[j] ||
(result.seg_offsets_[j] >= 0 &&
result.seg_offsets_[j] < 0);
});
}
for (auto i : idx) {
internal/core/src/exec/operator/groupby/SearchGroupByOperator.cpp
+15
View File
@@ -97,6 +97,21 @@ SearchGroupBy(const std::vector<std::shared_ptr<VectorIterator>>& iterators,
topk_per_nq_prefix_sum);
break;
}
case DataType::TIMESTAMPTZ: {
auto dataGetter =
GetDataGetter<int64_t>(segment, group_by_field_id);
GroupIteratorsByType<int64_t>(iterators,
search_info.topk_,
search_info.group_size_,
search_info.strict_group_size_,
*dataGetter,
group_by_values,
seg_offsets,
distances,
search_info.metric_type_,
topk_per_nq_prefix_sum);
break;
}
case DataType::BOOL: {
auto dataGetter = GetDataGetter<bool>(segment, group_by_field_id);
GroupIteratorsByType<bool>(iterators,
internal/core/src/plan/PlanNode.h
+3 -1
View File
@@ -445,7 +445,9 @@ class RescoresNode : public PlanNode {
const PlanNodeId& id,
const std::vector<std::shared_ptr<rescores::Scorer>> scorers,
const std::vector<PlanNodePtr>& sources = std::vector<PlanNodePtr>{})
: PlanNode(id), scorers_(std::move(scorers)), sources_{std::move(sources)} {
: PlanNode(id),
scorers_(std::move(scorers)),
sources_{std::move(sources)} {
}
DataType
internal/core/src/query/PlanProto.cpp
+5 -4
View File
@@ -211,12 +211,12 @@ ProtoParser::PlanNodeFromProto(const planpb::PlanNode& plan_node_proto) {
}
// if has score function, run filter and scorer at last
if (plan_node_proto.scorers_size() > 0){
if (plan_node_proto.scorers_size() > 0) {
std::vector<std::shared_ptr<rescores::Scorer>> scorers;
for (const auto& function: plan_node_proto.scorers()){
for (const auto& function : plan_node_proto.scorers()) {
scorers.push_back(ParseScorer(function));
}
plannode = std::make_shared<milvus::plan::RescoresNode>(
milvus::plan::GetNextPlanNodeId(), std::move(scorers), sources);
sources = std::vector<milvus::plan::PlanNodePtr>{plannode};
@@ -604,7 +604,8 @@ ProtoParser::ParseExprs(const proto::plan::Expr& expr_pb,
ExprInvalid, "expr type check failed, actual type: {}", result->type());
}
std::shared_ptr<rescores::Scorer> ProtoParser::ParseScorer(const proto::plan::ScoreFunction& function){
std::shared_ptr<rescores::Scorer>
ProtoParser::ParseScorer(const proto::plan::ScoreFunction& function) {
auto expr = ParseExprs(function.filter());
return std::make_shared<rescores::WeightScorer>(expr, function.weight());
}
internal/core/src/query/PlanProto.h
+1 -1
View File
@@ -54,7 +54,7 @@ class ProtoParser {
ParseExprs(const proto::plan::Expr& expr_pb,
TypeCheckFunction type_check = TypeIsBool);
std::shared_ptr<rescores::Scorer>
std::shared_ptr<rescores::Scorer>
ParseScorer(const proto::plan::ScoreFunction& function);
private:
internal/core/src/rescores/Scorer.h
+1 -1
View File
@@ -47,7 +47,7 @@ class WeightScorer : public Scorer {
}
float
weight() override{
weight() override {
return weight_;
}
internal/core/src/segcore/ChunkedSegmentSealedImpl.cpp
+10
View File
@@ -1705,6 +1705,16 @@ ChunkedSegmentSealedImpl::get_raw_data(FieldId field_id,
->mutable_data());
break;
}
case DataType::TIMESTAMPTZ: {
bulk_subscript_impl<int64_t>(column.get(),
seg_offsets,
count,
ret->mutable_scalars()
->mutable_timestamptz_data()
->mutable_data()
->mutable_data());
break;
}
case DataType::VECTOR_FLOAT: {
bulk_subscript_impl(field_meta.get_sizeof(),
column.get(),
internal/core/src/segcore/ConcurrentVector.cpp
+5
View File
@@ -93,6 +93,11 @@ VectorBase::set_data_raw(ssize_t element_offset,
return set_data_raw(
element_offset, FIELD_DATA(data, double).data(), element_count);
}
case DataType::TIMESTAMPTZ: {
return set_data_raw(element_offset,
FIELD_DATA(data, timestamptz).data(),
element_count);
}
case DataType::VARCHAR:
case DataType::TEXT: {
auto& field_data = FIELD_DATA(data, string);
internal/core/src/segcore/FieldIndexing.cpp
+3
View File
@@ -370,6 +370,9 @@ CreateIndex(const FieldMeta& field_meta,
case DataType::DOUBLE:
return std::make_unique<ScalarFieldIndexing<double>>(
field_meta, segcore_config);
case DataType::TIMESTAMPTZ:
return std::make_unique<ScalarFieldIndexing<int64_t>>(
field_meta, segcore_config);
case DataType::VARCHAR:
return std::make_unique<ScalarFieldIndexing<std::string>>(
field_meta, segcore_config);
internal/core/src/segcore/InsertRecord.h
+20 -12
View File
@@ -25,6 +25,7 @@
#include "common/EasyAssert.h"
#include "common/Schema.h"
#include "common/Types.h"
#include "log/Log.h"
#include "mmap/ChunkedColumn.h"
#include "segcore/AckResponder.h"
#include "segcore/ConcurrentVector.h"
@@ -270,49 +271,51 @@ class OffsetOrderedArray : public OffsetMap {
BitsetTypeView& bitset,
Condition condition) const override {
check_search();
auto lower_bound_comp = [](const std::pair<T, int64_t>& elem, const T& value) {
auto lower_bound_comp = [](const std::pair<T, int64_t>& elem,
const T& value) {
return elem.first < value;
};
auto upper_bound_comp = [](const T& value, const std::pair<T, int64_t>& elem) {
auto upper_bound_comp = [](const T& value,
const std::pair<T, int64_t>& elem) {
return value < elem.first;
};
const T& target = std::get<T>(pk);
if (op == proto::plan::OpType::Equal) {
auto it =
std::lower_bound(array_.begin(), array_.end(), target, lower_bound_comp);
auto it = std::lower_bound(
array_.begin(), array_.end(), target, lower_bound_comp);
for (; it != array_.end() && it->first == target; ++it) {
if (condition(it->second)) {
bitset[it->second] = true;
}
}
} else if (op == proto::plan::OpType::GreaterEqual) {
auto it =
std::lower_bound(array_.begin(), array_.end(), target, lower_bound_comp);
auto it = std::lower_bound(
array_.begin(), array_.end(), target, lower_bound_comp);
for (; it < array_.end(); ++it) {
if (condition(it->second)) {
bitset[it->second] = true;
}
}
} else if (op == proto::plan::OpType::GreaterThan) {
auto it =
std::upper_bound(array_.begin(), array_.end(), target, upper_bound_comp);
auto it = std::upper_bound(
array_.begin(), array_.end(), target, upper_bound_comp);
for (; it < array_.end(); ++it) {
if (condition(it->second)) {
bitset[it->second] = true;
}
}
} else if (op == proto::plan::OpType::LessEqual) {
auto it =
std::upper_bound(array_.begin(), array_.end(), target, upper_bound_comp);
auto it = std::upper_bound(
array_.begin(), array_.end(), target, upper_bound_comp);
for (auto ptr = array_.begin(); ptr < it; ++ptr) {
if (condition(ptr->second)) {
bitset[ptr->second] = true;
}
}
} else if (op == proto::plan::OpType::LessThan) {
auto it =
std::lower_bound(array_.begin(), array_.end(), target, lower_bound_comp);
auto it = std::lower_bound(
array_.begin(), array_.end(), target, lower_bound_comp);
for (auto ptr = array_.begin(); ptr < it; ++ptr) {
if (condition(ptr->second)) {
bitset[ptr->second] = true;
@@ -757,6 +760,11 @@ struct InsertRecord<false> : public InsertRecord<true> {
field_id, size_per_chunk, scalar_mmap_descriptor);
return;
}
case DataType::TIMESTAMPTZ: {
this->append_data<int64_t>(
field_id, size_per_chunk, scalar_mmap_descriptor);
return;
}
case DataType::VARCHAR:
case DataType::TEXT: {
this->append_data<std::string>(
internal/core/src/segcore/ReduceUtils.cpp
+15
View File
@@ -93,6 +93,21 @@ AssembleGroupByValues(
}
break;
}
case DataType::TIMESTAMPTZ: {
auto field_data =
group_by_values_field->mutable_timestamptz_data();
field_data->mutable_data()->Resize(group_by_val_size, 0);
for (std::size_t idx = 0; idx < group_by_val_size; idx++) {
if (group_by_vals[idx].has_value()) {
int64_t val =
std::get<int64_t>(group_by_vals[idx].value());
field_data->mutable_data()->Set(idx, val);
} else {
valid_data->Set(idx, false);
}
}
break;
}
case DataType::BOOL: {
auto field_data = group_by_values_field->mutable_bool_data();
field_data->mutable_data()->Resize(group_by_val_size, 0);
internal/core/src/segcore/SegmentChunkReader.cpp
+2
View File
@@ -175,6 +175,7 @@ SegmentChunkReader::GetChunkDataAccessor(DataType data_type,
return GetChunkDataAccessor<int32_t>(
field_id, index, current_chunk_id, current_chunk_pos);
case DataType::INT64:
case DataType::TIMESTAMPTZ:
return GetChunkDataAccessor<int64_t>(
field_id, index, current_chunk_id, current_chunk_pos);
case DataType::FLOAT:
@@ -286,6 +287,7 @@ SegmentChunkReader::GetChunkDataAccessor(DataType data_type,
case DataType::INT32:
return GetChunkDataAccessor<int32_t>(
field_id, chunk_id, data_barrier);
case DataType::TIMESTAMPTZ:
case DataType::INT64:
return GetChunkDataAccessor<int64_t>(
field_id, chunk_id, data_barrier);
internal/core/src/segcore/SegmentGrowingImpl.cpp
+10
View File
@@ -895,6 +895,16 @@ SegmentGrowingImpl::bulk_subscript(FieldId field_id,
->mutable_data());
break;
}
case DataType::TIMESTAMPTZ: {
bulk_subscript_impl<int64_t>(vec_ptr,
seg_offsets,
count,
result->mutable_scalars()
->mutable_timestamptz_data()
->mutable_data()
->mutable_data());
break;
}
case DataType::VARCHAR:
case DataType::TEXT: {
bulk_subscript_ptr_impl<std::string>(vec_ptr,
internal/core/src/segcore/SegmentInterface.cpp
+12
View File
@@ -501,6 +501,18 @@ SegmentInternalInterface::bulk_subscript_not_exist_field(
}
break;
}
case DataType::TIMESTAMPTZ: {
auto data_ptr = result->mutable_scalars()
->mutable_timestamptz_data()
->mutable_data()
->mutable_data();
for (int64_t i = 0; i < count; ++i) {
data_ptr[i] =
field_meta.default_value()->timestamptz_data();
}
break;
}
case DataType::VARCHAR: {
auto data_ptr = result->mutable_scalars()
->mutable_string_data()
internal/core/src/segcore/Utils.cpp
+46
View File
@@ -188,6 +188,14 @@ GetRawDataSizeOfDataArray(const DataArray* data,
}
break;
}
case DataType::TIMESTAMPTZ: {
for (auto& array_bytes : array_data) {
result +=
array_bytes.timestamptz_data().data_size() *
sizeof(int64_t);
}
break;
}
case DataType::VARCHAR:
case DataType::STRING:
case DataType::TEXT: {
@@ -295,6 +303,11 @@ CreateEmptyScalarDataArray(int64_t count, const FieldMeta& field_meta) {
obj->mutable_data()->Resize(count, 0);
break;
}
case DataType::TIMESTAMPTZ: {
auto obj = scalar_array->mutable_timestamptz_data();
obj->mutable_data()->Resize(count, 0);
break;
}
case DataType::VARCHAR:
case DataType::STRING:
case DataType::TEXT: {
@@ -461,6 +474,12 @@ CreateScalarDataArrayFrom(const void* data_raw,
obj->mutable_data()->Add(data, data + count);
break;
}
case DataType::TIMESTAMPTZ: {
auto data = reinterpret_cast<const int64_t*>(data_raw);
auto obj = scalar_array->mutable_timestamptz_data();
obj->mutable_data()->Add(data, data + count);
break;
}
case DataType::VARCHAR:
case DataType::TEXT: {
auto data = reinterpret_cast<const std::string*>(data_raw);
@@ -720,6 +739,13 @@ MergeDataArray(std::vector<MergeBase>& merge_bases,
*(obj->mutable_data()->Add()) = data[src_offset];
break;
}
case DataType::TIMESTAMPTZ: {
auto data = FIELD_DATA(src_field_data, timestamptz)
.data(); //Here is a marco
auto obj = scalar_array->mutable_timestamptz_data();
*(obj->mutable_data()->Add()) = data[src_offset];
break;
}
case DataType::VARCHAR:
case DataType::TEXT: {
auto& data = FIELD_DATA(src_field_data, string);
@@ -909,6 +935,26 @@ ReverseDataFromIndex(const index::IndexBase* index,
*(obj->mutable_data()) = {raw_data.begin(), raw_data.end()};
break;
}
case DataType::TIMESTAMPTZ: {
using IndexType = index::ScalarIndex<int64_t>;
auto ptr = dynamic_cast<const IndexType*>(index);
std::vector<int64_t> raw_data(count);
for (int64_t i = 0; i < count; ++i) {
auto raw = ptr->Reverse_Lookup(seg_offsets[i]);
// if has no value, means nullable must be true, no need to check nullable again
if (!raw.has_value()) {
valid_data[i] = false;
continue;
}
if (nullable) {
valid_data[i] = true;
}
raw_data[i] = raw.value();
auto obj = scalar_array->mutable_timestamptz_data();
*(obj->mutable_data()) = {raw_data.begin(), raw_data.end()};
break;
}
}
case DataType::VARCHAR: {
using IndexType = index::ScalarIndex<std::string>;
auto ptr = dynamic_cast<const IndexType*>(index);
internal/core/src/segcore/packed_writer_c.cpp
+8 -4
View File
@@ -149,8 +149,10 @@ WriteRecordBatch(CPackedWriter c_packed_writer,
auto import_schema = arrow::ImportSchema(schema);
if (!import_schema.ok()) {
return milvus::FailureCStatus(milvus::ErrorCode::FileWriteFailed,
"Failed to import schema: " + import_schema.status().ToString());
return milvus::FailureCStatus(
milvus::ErrorCode::FileWriteFailed,
"Failed to import schema: " +
import_schema.status().ToString());
}
auto arrow_schema = import_schema.ValueOrDie();
@@ -161,8 +163,10 @@ WriteRecordBatch(CPackedWriter c_packed_writer,
for (int i = 0; i < num_fields; i++) {
auto array = arrow::ImportArray(&arrays[i], &array_schemas[i]);
if (!array.ok()) {
return milvus::FailureCStatus(milvus::ErrorCode::FileWriteFailed,
"Failed to import array " + std::to_string(i) + ": " + array.status().ToString());
return milvus::FailureCStatus(
milvus::ErrorCode::FileWriteFailed,
"Failed to import array " + std::to_string(i) + ": " +
array.status().ToString());
}
all_arrays.push_back(array.ValueOrDie());
}
internal/core/src/storage/Util.cpp
+24
View File
@@ -198,6 +198,17 @@ AddPayloadToArrowBuilder(std::shared_ptr<arrow::ArrayBuilder> builder,
builder, double_data, payload.valid_data, nullable, length);
break;
}
case DataType::TIMESTAMPTZ: {
auto timestamptz_data = reinterpret_cast<int64_t*>(raw_data);
add_numeric_payload<int64_t, arrow::Int64Builder>(
builder,
timestamptz_data,
payload.valid_data,
nullable,
length);
break;
}
case DataType::VECTOR_FLOAT16:
case DataType::VECTOR_BFLOAT16:
case DataType::VECTOR_BINARY:
@@ -276,6 +287,9 @@ CreateArrowBuilder(DataType data_type) {
case DataType::DOUBLE: {
return std::make_shared<arrow::DoubleBuilder>();
}
case DataType::TIMESTAMPTZ: {
return std::make_shared<arrow::Int64Builder>();
}
case DataType::VARCHAR:
case DataType::STRING:
case DataType::TEXT: {
@@ -360,6 +374,9 @@ CreateArrowScalarFromDefaultValue(const FieldMeta& field_meta) {
case DataType::DOUBLE:
return std::make_shared<arrow::DoubleScalar>(
default_value.double_data());
case DataType::TIMESTAMPTZ:
return std::make_shared<arrow::Int64Scalar>(
default_value.timestamptz_data());
case DataType::VARCHAR:
case DataType::STRING:
case DataType::TEXT:
@@ -403,6 +420,10 @@ CreateArrowSchema(DataType data_type, bool nullable) {
return arrow::schema(
{arrow::field("val", arrow::float64(), nullable)});
}
case DataType::TIMESTAMPTZ: {
return arrow::schema(
{arrow::field("val", arrow::int64(), nullable)});
}
case DataType::VARCHAR:
case DataType::STRING:
case DataType::TEXT: {
@@ -948,6 +969,9 @@ CreateFieldData(const DataType& type,
case DataType::DOUBLE:
return std::make_shared<FieldData<double>>(
type, nullable, total_num_rows);
case DataType::TIMESTAMPTZ:
return std::make_shared<FieldData<int64_t>>(
type, nullable, total_num_rows);
case DataType::STRING:
case DataType::VARCHAR:
case DataType::TEXT:
internal/core/unittest/test_bitset.cpp
+22 -11
View File
@@ -503,7 +503,8 @@ TYPED_TEST_P(FindSuite, Neon) {
}
TYPED_TEST_P(FindSuite, Sve) {
#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && defined(BITSET_ENABLE_SVE_SUPPORT)
#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && \
defined(BITSET_ENABLE_SVE_SUPPORT)
using namespace milvus::bitset::detail::arm;
using impl_traits =
@@ -688,7 +689,8 @@ TYPED_TEST_P(InplaceCompareColumnSuite, Neon) {
//
TYPED_TEST_P(InplaceCompareColumnSuite, Sve) {
#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && defined(BITSET_ENABLE_SVE_SUPPORT)
#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && \
defined(BITSET_ENABLE_SVE_SUPPORT)
using namespace milvus::bitset::detail::arm;
using impl_traits =
@@ -878,7 +880,8 @@ TYPED_TEST_P(InplaceCompareValSuite, Neon) {
}
TYPED_TEST_P(InplaceCompareValSuite, Sve) {
#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && defined(BITSET_ENABLE_SVE_SUPPORT)
#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && \
defined(BITSET_ENABLE_SVE_SUPPORT)
using namespace milvus::bitset::detail::arm;
using impl_traits =
@@ -1073,7 +1076,8 @@ TYPED_TEST_P(InplaceWithinRangeColumnSuite, Neon) {
}
TYPED_TEST_P(InplaceWithinRangeColumnSuite, Sve) {
#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && defined(BITSET_ENABLE_SVE_SUPPORT)
#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && \
defined(BITSET_ENABLE_SVE_SUPPORT)
using namespace milvus::bitset::detail::arm;
using impl_traits =
@@ -1259,7 +1263,8 @@ TYPED_TEST_P(InplaceWithinRangeValSuite, Neon) {
}
TYPED_TEST_P(InplaceWithinRangeValSuite, Sve) {
#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && defined(BITSET_ENABLE_SVE_SUPPORT)
#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && \
defined(BITSET_ENABLE_SVE_SUPPORT)
using namespace milvus::bitset::detail::arm;
using impl_traits =
@@ -1712,7 +1717,8 @@ TYPED_TEST_P(InplaceArithCompareSuite, Neon) {
}
TYPED_TEST_P(InplaceArithCompareSuite, Sve) {
#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && defined(BITSET_ENABLE_SVE_SUPPORT)
#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && \
defined(BITSET_ENABLE_SVE_SUPPORT)
using namespace milvus::bitset::detail::arm;
using impl_traits =
@@ -1897,7 +1903,8 @@ TYPED_TEST_P(AppendSuite, Neon) {
}
TYPED_TEST_P(AppendSuite, Sve) {
#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && defined(BITSET_ENABLE_SVE_SUPPORT)
#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && \
defined(BITSET_ENABLE_SVE_SUPPORT)
using namespace milvus::bitset::detail::arm;
using impl_traits =
@@ -2062,7 +2069,8 @@ TYPED_TEST_P(CountSuite, Neon) {
}
TYPED_TEST_P(CountSuite, Sve) {
#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && defined(BITSET_ENABLE_SVE_SUPPORT)
#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && \
defined(BITSET_ENABLE_SVE_SUPPORT)
using namespace milvus::bitset::detail::arm;
using impl_traits =
@@ -2408,7 +2416,8 @@ TYPED_TEST_P(InplaceOpSuite, Neon) {
}
TYPED_TEST_P(InplaceOpSuite, Sve) {
#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && defined(BITSET_ENABLE_SVE_SUPPORT)
#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && \
defined(BITSET_ENABLE_SVE_SUPPORT)
using namespace milvus::bitset::detail::arm;
using impl_traits =
@@ -2636,7 +2645,8 @@ TYPED_TEST_P(InplaceOpMultipleSuite, Neon) {
}
TYPED_TEST_P(InplaceOpMultipleSuite, Sve) {
#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && defined(BITSET_ENABLE_SVE_SUPPORT)
#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && \
defined(BITSET_ENABLE_SVE_SUPPORT)
using namespace milvus::bitset::detail::arm;
using impl_traits =
@@ -2851,7 +2861,8 @@ TYPED_TEST_P(FillSuite, Neon) {
}
TYPED_TEST_P(FillSuite, Sve) {
#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && defined(BITSET_ENABLE_SVE_SUPPORT)
#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && \
defined(BITSET_ENABLE_SVE_SUPPORT)
using namespace milvus::bitset::detail::arm;
using impl_traits =
internal/core/unittest/test_c_api.cpp
+5
View File
@@ -203,6 +203,11 @@ generate_collection_schema(std::string metric_type, int dim) {
other_field_schema2->set_fieldid(102);
other_field_schema2->set_data_type(schema::DataType::Double);
auto other_field_schema3 = collection_schema.add_fields();
other_field_schema3->set_name("timestamptzField");
other_field_schema3->set_fieldid(103);
other_field_schema3->set_data_type(schema::DataType::Timestamptz);
std::string schema_string;
auto marshal = google::protobuf::TextFormat::PrintToString(
collection_schema, &schema_string);
internal/core/unittest/test_chunk.cpp
+119 -52
View File
@@ -76,6 +76,48 @@ TEST(chunk, test_int64_field) {
}
}
TEST(chunk, test_timestmamptz_field) {
FixedVector<int64_t> data = {
1, 2, 3, 4, 5}; // Timestamptz is stored as int64
auto field_data =
milvus::storage::CreateFieldData(storage::DataType::TIMESTAMPTZ);
field_data->FillFieldData(data.data(), data.size());
storage::InsertEventData event_data;
auto payload_reader =
std::make_shared<milvus::storage::PayloadReader>(field_data);
event_data.payload_reader = payload_reader;
auto ser_data = event_data.Serialize();
auto buffer = std::make_shared<arrow::io::BufferReader>(
ser_data.data() + 2 * sizeof(milvus::Timestamp),
ser_data.size() - 2 * sizeof(milvus::Timestamp));
parquet::arrow::FileReaderBuilder reader_builder;
auto s = reader_builder.Open(buffer);
EXPECT_TRUE(s.ok());
std::unique_ptr<parquet::arrow::FileReader> arrow_reader;
s = reader_builder.Build(&arrow_reader);
EXPECT_TRUE(s.ok());
std::shared_ptr<::arrow::RecordBatchReader> rb_reader;
s = arrow_reader->GetRecordBatchReader(&rb_reader);
EXPECT_TRUE(s.ok());
FieldMeta field_meta(FieldName("a"),
milvus::FieldId(1),
DataType::TIMESTAMPTZ,
false,
std::nullopt);
arrow::ArrayVector array_vec = read_single_column_batches(rb_reader);
auto chunk = create_chunk(field_meta, array_vec);
auto fixed_chunk = static_cast<FixedWidthChunk*>(chunk.get());
auto span = fixed_chunk->Span();
EXPECT_EQ(span.row_count(), data.size());
for (size_t i = 0; i < data.size(); ++i) {
auto n = *(int64_t*)((char*)span.data() + i * span.element_sizeof());
EXPECT_EQ(n, data[i]);
}
}
TEST(chunk, test_variable_field) {
FixedVector<std::string> data = {
"test1", "test2", "test3", "test4", "test5"};
@@ -614,16 +656,24 @@ TEST(chunk, test_sparse_float) {
TEST(chunk, test_lower_bound_string) {
// Test data: sorted strings
FixedVector<std::string> data = {
"apple", "banana", "cherry", "date", "elderberry",
"fig", "grape", "honeydew", "kiwi", "lemon"
};
auto field_data = milvus::storage::CreateFieldData(storage::DataType::VARCHAR);
FixedVector<std::string> data = {"apple",
"banana",
"cherry",
"date",
"elderberry",
"fig",
"grape",
"honeydew",
"kiwi",
"lemon"};
auto field_data =
milvus::storage::CreateFieldData(storage::DataType::VARCHAR);
field_data->FillFieldData(data.data(), data.size());
storage::InsertEventData event_data;
auto payload_reader = std::make_shared<milvus::storage::PayloadReader>(field_data);
auto payload_reader =
std::make_shared<milvus::storage::PayloadReader>(field_data);
event_data.payload_reader = payload_reader;
auto ser_data = event_data.Serialize();
auto buffer = std::make_shared<arrow::io::BufferReader>(
@@ -653,27 +703,28 @@ TEST(chunk, test_lower_bound_string) {
// Test cases for lower_bound_string
// Case 1: Target exists in the middle
EXPECT_EQ(string_chunk->lower_bound_string("cherry"), 2);
// Case 2: Target exists at the beginning
EXPECT_EQ(string_chunk->lower_bound_string("apple"), 0);
// Case 3: Target exists at the end
EXPECT_EQ(string_chunk->lower_bound_string("lemon"), 9);
// Case 4: Target doesn't exist, should return insertion point
EXPECT_EQ(string_chunk->lower_bound_string("blueberry"), 2); // between banana and cherry
EXPECT_EQ(string_chunk->lower_bound_string("mango"), 10); // after lemon
EXPECT_EQ(string_chunk->lower_bound_string("apricot"), 1); // after apple
EXPECT_EQ(string_chunk->lower_bound_string("blueberry"),
2); // between banana and cherry
EXPECT_EQ(string_chunk->lower_bound_string("mango"), 10); // after lemon
EXPECT_EQ(string_chunk->lower_bound_string("apricot"), 1); // after apple
// Case 5: Target is less than all elements
EXPECT_EQ(string_chunk->lower_bound_string("aardvark"), 0);
// Case 6: Target is greater than all elements
EXPECT_EQ(string_chunk->lower_bound_string("zebra"), 10);
// Case 7: Empty string edge case
EXPECT_EQ(string_chunk->lower_bound_string(""), 0);
// Case 8: Duplicate elements (if they existed, lower_bound would return first occurrence)
// Since our data has no duplicates, test with existing elements
EXPECT_EQ(string_chunk->lower_bound_string("banana"), 1);
@@ -681,16 +732,24 @@ TEST(chunk, test_lower_bound_string) {
TEST(chunk, test_upper_bound_string) {
// Test data: sorted strings
FixedVector<std::string> data = {
"apple", "banana", "cherry", "date", "elderberry",
"fig", "grape", "honeydew", "kiwi", "lemon"
};
auto field_data = milvus::storage::CreateFieldData(storage::DataType::VARCHAR);
FixedVector<std::string> data = {"apple",
"banana",
"cherry",
"date",
"elderberry",
"fig",
"grape",
"honeydew",
"kiwi",
"lemon"};
auto field_data =
milvus::storage::CreateFieldData(storage::DataType::VARCHAR);
field_data->FillFieldData(data.data(), data.size());
storage::InsertEventData event_data;
auto payload_reader = std::make_shared<milvus::storage::PayloadReader>(field_data);
auto payload_reader =
std::make_shared<milvus::storage::PayloadReader>(field_data);
event_data.payload_reader = payload_reader;
auto ser_data = event_data.Serialize();
auto buffer = std::make_shared<arrow::io::BufferReader>(
@@ -719,28 +778,31 @@ TEST(chunk, test_upper_bound_string) {
// Test cases for upper_bound_string
// Case 1: Target exists in the middle
EXPECT_EQ(string_chunk->upper_bound_string("cherry"), 3); // points to next element
EXPECT_EQ(string_chunk->upper_bound_string("cherry"),
3); // points to next element
// Case 2: Target exists at the beginning
EXPECT_EQ(string_chunk->upper_bound_string("apple"), 1); // points to next element
EXPECT_EQ(string_chunk->upper_bound_string("apple"),
1); // points to next element
// Case 3: Target exists at the end
EXPECT_EQ(string_chunk->upper_bound_string("lemon"), 10); // points to end
// Case 4: Target doesn't exist, should return insertion point
EXPECT_EQ(string_chunk->upper_bound_string("blueberry"), 2); // between banana and cherry
EXPECT_EQ(string_chunk->upper_bound_string("mango"), 10); // after lemon
EXPECT_EQ(string_chunk->upper_bound_string("apricot"), 1); // after apple
EXPECT_EQ(string_chunk->upper_bound_string("blueberry"),
2); // between banana and cherry
EXPECT_EQ(string_chunk->upper_bound_string("mango"), 10); // after lemon
EXPECT_EQ(string_chunk->upper_bound_string("apricot"), 1); // after apple
// Case 5: Target is less than all elements
EXPECT_EQ(string_chunk->upper_bound_string("aardvark"), 0);
// Case 6: Target is greater than all elements
EXPECT_EQ(string_chunk->upper_bound_string("zebra"), 10);
// Case 7: Empty string edge case
EXPECT_EQ(string_chunk->upper_bound_string(""), 0);
// Case 8: Test with existing elements (upper_bound points to next element)
EXPECT_EQ(string_chunk->upper_bound_string("banana"), 2);
EXPECT_EQ(string_chunk->upper_bound_string("date"), 4);
@@ -749,12 +811,14 @@ TEST(chunk, test_upper_bound_string) {
TEST(chunk, test_binary_search_methods_edge_cases) {
// Test with single element
FixedVector<std::string> single_data = {"middle"};
auto field_data = milvus::storage::CreateFieldData(storage::DataType::VARCHAR);
auto field_data =
milvus::storage::CreateFieldData(storage::DataType::VARCHAR);
field_data->FillFieldData(single_data.data(), single_data.size());
storage::InsertEventData event_data;
auto payload_reader = std::make_shared<milvus::storage::PayloadReader>(field_data);
auto payload_reader =
std::make_shared<milvus::storage::PayloadReader>(field_data);
event_data.payload_reader = payload_reader;
auto ser_data = event_data.Serialize();
auto buffer = std::make_shared<arrow::io::BufferReader>(
@@ -793,14 +857,15 @@ TEST(chunk, test_binary_search_methods_edge_cases) {
TEST(chunk, test_binary_search_methods_duplicates) {
// Test with duplicate elements (if the data had duplicates)
FixedVector<std::string> data = {
"apple", "apple", "banana", "banana", "banana", "cherry"
};
auto field_data = milvus::storage::CreateFieldData(storage::DataType::VARCHAR);
"apple", "apple", "banana", "banana", "banana", "cherry"};
auto field_data =
milvus::storage::CreateFieldData(storage::DataType::VARCHAR);
field_data->FillFieldData(data.data(), data.size());
storage::InsertEventData event_data;
auto payload_reader = std::make_shared<milvus::storage::PayloadReader>(field_data);
auto payload_reader =
std::make_shared<milvus::storage::PayloadReader>(field_data);
event_data.payload_reader = payload_reader;
auto ser_data = event_data.Serialize();
auto buffer = std::make_shared<arrow::io::BufferReader>(
@@ -832,7 +897,7 @@ TEST(chunk, test_binary_search_methods_duplicates) {
EXPECT_EQ(string_chunk->lower_bound_string("apple"), 0);
EXPECT_EQ(string_chunk->lower_bound_string("banana"), 2);
EXPECT_EQ(string_chunk->lower_bound_string("cherry"), 5);
// upper_bound should return the position after the last occurrence
EXPECT_EQ(string_chunk->upper_bound_string("apple"), 2);
EXPECT_EQ(string_chunk->upper_bound_string("banana"), 5);
@@ -842,14 +907,15 @@ TEST(chunk, test_binary_search_methods_duplicates) {
TEST(chunk, test_binary_search_methods_comparison) {
// Test to verify the relationship between lower_bound and upper_bound
FixedVector<std::string> data = {
"a", "b", "c", "d", "e", "f", "g", "h", "i", "j"
};
auto field_data = milvus::storage::CreateFieldData(storage::DataType::VARCHAR);
"a", "b", "c", "d", "e", "f", "g", "h", "i", "j"};
auto field_data =
milvus::storage::CreateFieldData(storage::DataType::VARCHAR);
field_data->FillFieldData(data.data(), data.size());
storage::InsertEventData event_data;
auto payload_reader = std::make_shared<milvus::storage::PayloadReader>(field_data);
auto payload_reader =
std::make_shared<milvus::storage::PayloadReader>(field_data);
event_data.payload_reader = payload_reader;
auto ser_data = event_data.Serialize();
auto buffer = std::make_shared<arrow::io::BufferReader>(
@@ -877,11 +943,12 @@ TEST(chunk, test_binary_search_methods_comparison) {
auto string_chunk = static_cast<StringChunk*>(chunk.get());
// Test that upper_bound >= lower_bound for any target
for (const auto& target : {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "z", ""}) {
for (const auto& target :
{"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "z", ""}) {
auto lb = string_chunk->lower_bound_string(target);
auto ub = string_chunk->upper_bound_string(target);
EXPECT_GE(ub, lb) << "For target: " << target;
// For existing elements, upper_bound should be lower_bound + 1
if (std::find(data.begin(), data.end(), target) != data.end()) {
EXPECT_EQ(ub, lb + 1) << "For existing target: " << target;
internal/core/unittest/test_chunk_vector.cpp
+6
View File
@@ -50,6 +50,8 @@ TEST_F(ChunkVectorTest, FillDataWithMmap) {
auto int64_field = schema->AddDebugField("int64", DataType::INT64);
auto float_field = schema->AddDebugField("float", DataType::FLOAT);
auto double_field = schema->AddDebugField("double", DataType::DOUBLE);
auto timestamptz_field =
schema->AddDebugField("timestamptz", DataType::TIMESTAMPTZ);
auto varchar_field = schema->AddDebugField("varchar", DataType::VARCHAR);
auto json_field = schema->AddDebugField("json", DataType::JSON);
auto int_array_field =
@@ -115,6 +117,8 @@ TEST_F(ChunkVectorTest, FillDataWithMmap) {
float_field, ids_ds->GetIds(), num_inserted);
auto double_result = segment->bulk_subscript(
double_field, ids_ds->GetIds(), num_inserted);
auto timestamptz_result = segment->bulk_subscript(
timestamptz_field, ids_ds->GetIds(), num_inserted);
auto varchar_result = segment->bulk_subscript(
varchar_field, ids_ds->GetIds(), num_inserted);
auto json_result =
@@ -152,6 +156,8 @@ TEST_F(ChunkVectorTest, FillDataWithMmap) {
num_inserted);
EXPECT_EQ(double_result->scalars().double_data().data_size(),
num_inserted);
EXPECT_EQ(timestamptz_result->scalars().timestamptz_data().data_size(),
num_inserted);
EXPECT_EQ(varchar_result->scalars().string_data().data_size(),
num_inserted);
EXPECT_EQ(json_result->scalars().json_data().data_size(), num_inserted);
internal/core/unittest/test_chunked_segment.cpp
+25 -25
View File
@@ -441,51 +441,51 @@ TEST_P(TestChunkSegment, TestPkRange) {
// Test cases for sorted PK
BitsetType bitset_sorted(chunk_num * test_data_count);
BitsetTypeView bitset_sorted_view(bitset_sorted);
// Test Equal operation
if (pk_is_string) {
segment_impl->pk_range(proto::plan::OpType::Equal,
PkType("test1"),
Timestamp(99999),
bitset_sorted_view);
segment_impl->pk_range(proto::plan::OpType::Equal,
PkType("test1"),
Timestamp(99999),
bitset_sorted_view);
EXPECT_EQ(1, bitset_sorted_view.count());
} else {
segment_impl->pk_range(proto::plan::OpType::Equal,
PkType(1),
Timestamp(99999),
bitset_sorted_view);
segment_impl->pk_range(proto::plan::OpType::Equal,
PkType(1),
Timestamp(99999),
bitset_sorted_view);
EXPECT_EQ(1, bitset_sorted_view.count());
}
// Test LessEqual operation
bitset_sorted.reset();
if (pk_is_string) {
segment_impl->pk_range(proto::plan::OpType::LessEqual,
PkType("test100"),
Timestamp(99999),
bitset_sorted_view);
segment_impl->pk_range(proto::plan::OpType::LessEqual,
PkType("test100"),
Timestamp(99999),
bitset_sorted_view);
// only 'test0', 'test1', 'test10' are less than 'test100'
EXPECT_EQ(bitset_sorted_view.count(), 4);
} else {
segment_impl->pk_range(proto::plan::OpType::LessEqual,
PkType(100),
Timestamp(99999),
bitset_sorted_view);
segment_impl->pk_range(proto::plan::OpType::LessEqual,
PkType(100),
Timestamp(99999),
bitset_sorted_view);
EXPECT_EQ(bitset_sorted_view.count(), 101);
}
bitset_sorted.reset();
if (pk_is_string) {
segment_impl->pk_range(proto::plan::OpType::Equal,
PkType(std::string("non_existent_pk")),
Timestamp(99999),
bitset_sorted_view);
segment_impl->pk_range(proto::plan::OpType::Equal,
PkType(std::string("non_existent_pk")),
Timestamp(99999),
bitset_sorted_view);
EXPECT_EQ(0, bitset_sorted_view.count());
} else {
segment_impl->pk_range(proto::plan::OpType::Equal,
PkType(int64_t(999999)),
Timestamp(99999),
bitset_sorted_view);
segment_impl->pk_range(proto::plan::OpType::Equal,
PkType(int64_t(999999)),
Timestamp(99999),
bitset_sorted_view);
EXPECT_EQ(0, bitset_sorted_view.count());
}
}
internal/core/unittest/test_data_codec.cpp
+66
View File
@@ -553,6 +553,72 @@ TEST(storage, InsertDataDoubleNullable) {
delete[] valid_data;
}
TEST(storage, InsertDataTimestamptz) {
FixedVector<int64_t> data = {
1000000000, 2000000000, 3000000000, 400000, 5000};
auto field_data =
milvus::storage::CreateFieldData(storage::DataType::TIMESTAMPTZ, false);
field_data->FillFieldData(data.data(), data.size());
auto payload_reader =
std::make_shared<milvus::storage::PayloadReader>(field_data);
storage::InsertData insert_data(payload_reader);
storage::FieldDataMeta field_data_meta{100, 101, 102, 103};
insert_data.SetFieldDataMeta(field_data_meta);
insert_data.SetTimestamps(0, 100);
auto serialized_bytes = insert_data.Serialize(storage::StorageType::Remote);
std::shared_ptr<uint8_t[]> serialized_data_ptr(serialized_bytes.data(),
[&](uint8_t*) {});
auto new_insert_data = storage::DeserializeFileData(
serialized_data_ptr, serialized_bytes.size());
ASSERT_EQ(new_insert_data->GetCodecType(), storage::InsertDataType);
ASSERT_EQ(new_insert_data->GetTimeRage(),
std::make_pair(Timestamp(0), Timestamp(100)));
auto new_payload = new_insert_data->GetFieldData();
ASSERT_EQ(new_payload->get_data_type(), storage::DataType::TIMESTAMPTZ);
ASSERT_EQ(new_payload->get_num_rows(), data.size());
ASSERT_EQ(new_payload->get_null_count(), 0);
FixedVector<int64_t> new_data(data.size());
memcpy(new_data.data(), new_payload->Data(), new_payload->DataSize());
ASSERT_EQ(data, new_data);
}
TEST(storage, InsertDataTimestamptzNullable) {
FixedVector<int64_t> data = {
1000000000, 2000000000, 3000000000, 400000, 5000};
auto field_data =
milvus::storage::CreateFieldData(storage::DataType::TIMESTAMPTZ, true);
uint8_t* valid_data = new uint8_t[1]{0xF3};
field_data->FillFieldData(data.data(), valid_data, data.size(), 0);
auto payload_reader =
std::make_shared<milvus::storage::PayloadReader>(field_data);
storage::InsertData insert_data(payload_reader);
storage::FieldDataMeta field_data_meta{100, 101, 102, 103};
insert_data.SetFieldDataMeta(field_data_meta);
insert_data.SetTimestamps(0, 100);
auto serialized_bytes = insert_data.Serialize(storage::StorageType::Remote);
std::shared_ptr<uint8_t[]> serialized_data_ptr(serialized_bytes.data(),
[&](uint8_t*) {});
auto new_insert_data = storage::DeserializeFileData(
serialized_data_ptr, serialized_bytes.size());
ASSERT_EQ(new_insert_data->GetCodecType(), storage::InsertDataType);
ASSERT_EQ(new_insert_data->GetTimeRage(),
std::make_pair(Timestamp(0), Timestamp(100)));
auto new_payload = new_insert_data->GetFieldData();
ASSERT_EQ(new_payload->get_data_type(), storage::DataType::TIMESTAMPTZ);
ASSERT_EQ(new_payload->get_num_rows(), data.size());
FixedVector<int64_t> new_data(data.size());
memcpy(new_data.data(), new_payload->Data(), new_payload->DataSize());
data = {1000000000, 2000000000, 0, 0, 5000};
ASSERT_EQ(data, new_data);
ASSERT_EQ(new_payload->get_null_count(), 2);
ASSERT_EQ(*new_payload->ValidData(), *valid_data);
delete[] valid_data;
}
TEST(storage, InsertDataFloatVector) {
std::vector<float> data = {1, 2, 3, 4, 5, 6, 7, 8};
int DIM = 2;
internal/core/unittest/test_group_by.cpp
+51 -1
View File
@@ -63,12 +63,14 @@ TEST(GroupBY, SealedIndex) {
auto schema = std::make_shared<Schema>();
auto vec_fid = schema->AddDebugField(
"fakevec", DataType::VECTOR_FLOAT, dim, knowhere::metric::L2);
auto int8_fid = schema->AddDebugField("int8", DataType::INT8);
auto int8_fid = schema->AddDebugField("int8", DataType::INT8); // id 101
auto int16_fid = schema->AddDebugField("int16", DataType::INT16);
auto int32_fid = schema->AddDebugField("int32", DataType::INT32);
auto int64_fid = schema->AddDebugField("int64", DataType::INT64);
auto str_fid = schema->AddDebugField("string1", DataType::VARCHAR);
auto bool_fid = schema->AddDebugField("bool", DataType::BOOL);
auto timestamptz_fid =
schema->AddDebugField("timestamptz", DataType::TIMESTAMPTZ);
schema->set_primary_field_id(str_fid);
size_t N = 50;
@@ -393,6 +395,54 @@ TEST(GroupBY, SealedIndex) {
}
ASSERT_TRUE(bools_map.size() == 2); //bool values cannot exceed two
}
//10. search group by timestamptz
{
const char* raw_plan = R"(vector_anns: <
field_id: 100
query_info: <
topk: 100
metric_type: "L2"
search_params: "{\"ef\": 10}"
group_by_field_id: 107
group_size: 3
>
placeholder_tag: "$0"
>)";
proto::plan::PlanNode plan_node;
auto ok =
google::protobuf::TextFormat::ParseFromString(raw_plan, &plan_node);
auto plan = CreateSearchPlanFromPlanNode(schema, plan_node);
auto num_queries = 1;
auto seed = 1024;
auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, seed);
auto ph_group =
ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
auto search_result =
segment->Search(plan.get(), ph_group.get(), 1L << 63);
CheckGroupBySearchResult(*search_result, topK, num_queries, false);
auto& group_by_values = search_result->group_by_values_.value();
int size = group_by_values.size();
ASSERT_EQ(20, size);
//as the total data is 0,0,....6,6, so there will be 7 buckets with [3,3,3,3,3,3,2] items respectively
//so there will be 20 items returned
std::unordered_map<int64_t, int> timestamptz_map;
float lastDistance = 0.0;
for (size_t i = 0; i < size; i++) {
if (std::holds_alternative<int64_t>(group_by_values[i].value())) {
int16_t g_val = std::get<int64_t>(group_by_values[i].value());
timestamptz_map[g_val] += 1;
ASSERT_TRUE(timestamptz_map[g_val] <= group_size);
auto distance = search_result->distances_.at(i);
ASSERT_TRUE(lastDistance <= distance);
lastDistance = distance;
}
}
ASSERT_TRUE(timestamptz_map.size() == 7);
}
}
TEST(GroupBY, SealedData) {
internal/core/unittest/test_growing.cpp
+14
View File
@@ -144,6 +144,8 @@ TEST_P(GrowingTest, FillData) {
auto int16_field = schema->AddDebugField("int16", DataType::INT16);
auto int32_field = schema->AddDebugField("int32", DataType::INT32);
auto int64_field = schema->AddDebugField("int64", DataType::INT64);
auto timestamptz_field =
schema->AddDebugField("timestamptz", DataType::TIMESTAMPTZ);
auto float_field = schema->AddDebugField("float", DataType::FLOAT);
auto double_field = schema->AddDebugField("double", DataType::DOUBLE);
auto varchar_field = schema->AddDebugField("varchar", DataType::VARCHAR);
@@ -207,6 +209,8 @@ TEST_P(GrowingTest, FillData) {
float_field, ids_ds->GetIds(), num_inserted);
auto double_result = segment->bulk_subscript(
double_field, ids_ds->GetIds(), num_inserted);
auto timestamptz_result = segment->bulk_subscript(
timestamptz_field, ids_ds->GetIds(), num_inserted);
auto varchar_result = segment->bulk_subscript(
varchar_field, ids_ds->GetIds(), num_inserted);
auto json_result =
@@ -232,6 +236,8 @@ TEST_P(GrowingTest, FillData) {
EXPECT_EQ(int32_result->scalars().int_data().data_size(), num_inserted);
EXPECT_EQ(int64_result->scalars().long_data().data_size(),
num_inserted);
EXPECT_EQ(timestamptz_result->scalars().timestamptz_data().data_size(),
num_inserted);
EXPECT_EQ(float_result->scalars().float_data().data_size(),
num_inserted);
EXPECT_EQ(double_result->scalars().double_data().data_size(),
@@ -269,6 +275,7 @@ TEST_P(GrowingTest, FillData) {
EXPECT_EQ(int64_result->valid_data_size(), 0);
EXPECT_EQ(float_result->valid_data_size(), 0);
EXPECT_EQ(double_result->valid_data_size(), 0);
EXPECT_EQ(timestamptz_result->valid_data_size(), 0);
EXPECT_EQ(varchar_result->valid_data_size(), 0);
EXPECT_EQ(json_result->valid_data_size(), 0);
EXPECT_EQ(int_array_result->valid_data_size(), 0);
@@ -290,6 +297,8 @@ TEST(Growing, FillNullableData) {
auto int64_field = schema->AddDebugField("int64", DataType::INT64);
auto float_field = schema->AddDebugField("float", DataType::FLOAT, true);
auto double_field = schema->AddDebugField("double", DataType::DOUBLE, true);
auto timestamptz_field =
schema->AddDebugField("timestamptz", DataType::TIMESTAMPTZ, true);
auto varchar_field =
schema->AddDebugField("varchar", DataType::VARCHAR, true);
auto json_field = schema->AddDebugField("json", DataType::JSON, true);
@@ -353,6 +362,8 @@ TEST(Growing, FillNullableData) {
float_field, ids_ds->GetIds(), num_inserted);
auto double_result = segment->bulk_subscript(
double_field, ids_ds->GetIds(), num_inserted);
auto timestamptz_result = segment->bulk_subscript(
timestamptz_field, ids_ds->GetIds(), num_inserted);
auto varchar_result = segment->bulk_subscript(
varchar_field, ids_ds->GetIds(), num_inserted);
auto json_result =
@@ -382,6 +393,8 @@ TEST(Growing, FillNullableData) {
num_inserted);
EXPECT_EQ(double_result->scalars().double_data().data_size(),
num_inserted);
EXPECT_EQ(timestamptz_result->scalars().timestamptz_data().data_size(),
num_inserted);
EXPECT_EQ(varchar_result->scalars().string_data().data_size(),
num_inserted);
EXPECT_EQ(json_result->scalars().json_data().data_size(), num_inserted);
@@ -405,6 +418,7 @@ TEST(Growing, FillNullableData) {
EXPECT_EQ(int32_result->valid_data_size(), num_inserted);
EXPECT_EQ(float_result->valid_data_size(), num_inserted);
EXPECT_EQ(double_result->valid_data_size(), num_inserted);
EXPECT_EQ(timestamptz_result->valid_data_size(), num_inserted);
EXPECT_EQ(varchar_result->valid_data_size(), num_inserted);
EXPECT_EQ(json_result->valid_data_size(), num_inserted);
EXPECT_EQ(int_array_result->valid_data_size(), num_inserted);
internal/core/unittest/test_packed_c.cpp
+2 -1
View File
@@ -84,7 +84,8 @@ TEST(CPackedTest, PackedWriterAndReader) {
struct ArrowArray arrays[] = {carray};
struct ArrowSchema array_schemas[] = {cschema};
c_status = WriteRecordBatch(c_packed_writer, arrays, array_schemas, &c_origin_schema);
c_status = WriteRecordBatch(
c_packed_writer, arrays, array_schemas, &c_origin_schema);
EXPECT_EQ(c_status.error_code, 0);
c_status = CloseWriter(c_packed_writer);
internal/core/unittest/test_storage.cpp
+11
View File
@@ -171,6 +171,17 @@ TEST_F(StorageUtilTest, CreateArrowScalarFromDefaultValue) {
auto scalar = CreateArrowScalarFromDefaultValue(double_field);
ASSERT_TRUE(scalar->ApproxEquals(arrow::DoubleScalar(1.0f)));
}
{
DefaultValueType default_value;
default_value.set_timestamptz_data(123456789);
FieldMeta timestamptz_field(FieldName("f"),
FieldId(100),
DataType::TIMESTAMPTZ,
false,
default_value);
auto scalar = CreateArrowScalarFromDefaultValue(timestamptz_field);
ASSERT_TRUE(scalar->Equals(*arrow::MakeScalar(int64_t(123456789))));
}
{
DefaultValueType default_value;
default_value.set_bool_data(true);
internal/core/unittest/test_types.cpp
+5
View File
@@ -55,6 +55,11 @@ TEST(GetArrowDataTypeTest, DOUBLE) {
ASSERT_TRUE(result->Equals(arrow::float64()));
}
TEST(GetArrowDataTypeTest, TIMESTAMPTZ) {
auto result = GetArrowDataType(DataType::TIMESTAMPTZ);
ASSERT_TRUE(result->Equals(arrow::int64()));
}
TEST(GetArrowDataTypeTest, STRING_TYPES) {
auto result1 = GetArrowDataType(DataType::STRING);
auto result2 = GetArrowDataType(DataType::VARCHAR);
internal/core/unittest/test_utils/DataGen.h
+39
View File
@@ -12,6 +12,7 @@
#pragma once
#include <boost/algorithm/string/predicate.hpp>
#include <cstdint>
#include <cstring>
#include <memory>
#include <random>
@@ -225,6 +226,15 @@ struct GeneratedData {
std::copy_n(src_data, raw_->num_rows(), ret.data());
break;
}
case DataType::TIMESTAMPTZ: {
auto src_data = reinterpret_cast<const T*>(
target_field_data.scalars()
.timestamptz_data()
.data()
.data());
std::copy_n(src_data, raw_->num_rows(), ret.data());
break;
}
case DataType::VARCHAR: {
auto ret_data =
reinterpret_cast<std::string*>(ret.data());
@@ -367,6 +377,8 @@ inline SchemaPtr CreateTestSchema() {
auto int32_field =
schema->AddDebugField("int32", milvus::DataType::INT32, true);
auto int64_field = schema->AddDebugField("int64", milvus::DataType::INT64);
auto timestamptz_field =
schema->AddDebugField("timestamptz", DataType::TIMESTAMPTZ, true);
auto float_field =
schema->AddDebugField("float", milvus::DataType::FLOAT, true);
auto double_field =
@@ -638,6 +650,19 @@ DataGen(SchemaPtr schema,
insert_cols(data, N, field_meta, random_valid);
break;
}
case DataType::TIMESTAMPTZ: {
vector<int64_t> data(N);
for (int i = 0; i < N; ++i) {
int64_t x = 0;
if (random_val)
x = random() % (2 * N);
else
x = i / repeat_count;
data[i] = x;
}
insert_cols(data, N, field_meta, random_valid);
break;
}
case DataType::INT64: {
vector<int64_t> data(N);
for (int i = 0; i < N; i++) {
@@ -1304,6 +1329,18 @@ CreateFieldDataFromDataArray(ssize_t raw_count,
}
break;
}
case DataType::TIMESTAMPTZ: {
auto raw_data =
data->scalars().timestamptz_data().data().data();
if (field_meta.is_nullable()) {
auto raw_valid_data = data->valid_data().data();
createNullableFieldData(
raw_data, raw_valid_data, DataType::TIMESTAMPTZ, dim);
} else {
createFieldData(raw_data, DataType::TIMESTAMPTZ, dim);
}
break;
}
case DataType::VARCHAR: {
auto begin = data->scalars().string_data().data().begin();
auto end = data->scalars().string_data().data().end();
@@ -1605,6 +1642,8 @@ gen_all_data_types_schema() {
schema->AddDebugField("float", milvus::DataType::FLOAT, true);
auto double_field =
schema->AddDebugField("double", milvus::DataType::DOUBLE, true);
auto timestamptz_field = schema->AddDebugField(
"timestamptz", milvus::DataType::TIMESTAMPTZ, true);
auto varchar_field =
schema->AddDebugField("varchar", milvus::DataType::VARCHAR, true);
auto json_field =
internal/datanode/importv2/util.go
+8
View File
@@ -327,6 +327,14 @@ func AppendNullableDefaultFieldsData(schema *schemapb.CollectionSchema, data *st
} else if nullable {
err = appender.AppendNull(fieldData, rowNum)
}
case schemapb.DataType_Timestamptz:
appender := &nullDefaultAppender[int64]{}
if defaultVal != nil {
v := defaultVal.GetTimestamptzData()
err = appender.AppendDefault(fieldData, v, rowNum)
} else if nullable {
err = appender.AppendNull(fieldData, rowNum)
}
case schemapb.DataType_VarChar:
appender := &nullDefaultAppender[string]{}
if defaultVal != nil {
internal/distributed/proxy/httpserver/utils.go
+35 -1
View File
@@ -439,7 +439,7 @@ func checkAndSetData(body []byte, collSchema *schemapb.CollectionSchema, partial
return merr.WrapErrParameterInvalid(schemapb.DataType_name[int32(fieldType)], dataString, err.Error()), reallyDataArray, validDataMap
}
reallyData[fieldName] = result
case schemapb.DataType_Int64:
case schemapb.DataType_Int64, schemapb.DataType_Timestamptz:
result, err := json.Number(dataString).Int64()
if err != nil {
return merr.WrapErrParameterInvalid(schemapb.DataType_name[int32(fieldType)], dataString, err.Error()), reallyDataArray, validDataMap
@@ -775,6 +775,8 @@ func anyToColumns(rows []map[string]interface{}, validDataMap map[string][]bool,
data = make([]float32, 0, rowsLen)
case schemapb.DataType_Double:
data = make([]float64, 0, rowsLen)
case schemapb.DataType_Timestamptz:
data = make([]int64, 0, rowsLen)
case schemapb.DataType_String:
data = make([]string, 0, rowsLen)
case schemapb.DataType_VarChar:
@@ -871,6 +873,8 @@ func anyToColumns(rows []map[string]interface{}, validDataMap map[string][]bool,
nameColumns[field.Name] = append(nameColumns[field.Name].([]int64), candi.v.Interface().(int64))
case schemapb.DataType_Float:
nameColumns[field.Name] = append(nameColumns[field.Name].([]float32), candi.v.Interface().(float32))
case schemapb.DataType_Timestamptz:
nameColumns[field.Name] = append(nameColumns[field.Name].([]int64), candi.v.Interface().(int64))
case schemapb.DataType_Double:
nameColumns[field.Name] = append(nameColumns[field.Name].([]float64), candi.v.Interface().(float64))
case schemapb.DataType_String:
@@ -1023,6 +1027,16 @@ func anyToColumns(rows []map[string]interface{}, validDataMap map[string][]bool,
},
},
}
case schemapb.DataType_Timestamptz:
colData.Field = &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_TimestamptzData{
TimestamptzData: &schemapb.TimestamptzArray{
Data: column.([]int64),
},
},
},
}
case schemapb.DataType_String:
colData.Field = &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
@@ -1358,6 +1372,8 @@ func buildQueryResp(rowsNum int64, needFields []string, fieldDataList []*schemap
rowsNum = int64(len(fieldDataList[0].GetScalars().GetFloatData().GetData()))
case schemapb.DataType_Double:
rowsNum = int64(len(fieldDataList[0].GetScalars().GetDoubleData().GetData()))
case schemapb.DataType_Timestamptz:
rowsNum = int64(len(fieldDataList[0].GetScalars().GetTimestamptzData().GetData()))
case schemapb.DataType_String:
rowsNum = int64(len(fieldDataList[0].GetScalars().GetStringData().GetData()))
case schemapb.DataType_VarChar:
@@ -1461,6 +1477,12 @@ func buildQueryResp(rowsNum int64, needFields []string, fieldDataList []*schemap
continue
}
row[fieldDataList[j].GetFieldName()] = fieldDataList[j].GetScalars().GetDoubleData().GetData()[i]
case schemapb.DataType_Timestamptz:
if len(fieldDataList[j].GetValidData()) != 0 && !fieldDataList[j].GetValidData()[i] {
row[fieldDataList[j].GetFieldName()] = nil
continue
}
row[fieldDataList[j].FieldName] = fieldDataList[j].GetScalars().GetTimestamptzData().GetData()[i]
case schemapb.DataType_String:
if len(fieldDataList[j].GetValidData()) != 0 && !fieldDataList[j].GetValidData()[i] {
row[fieldDataList[j].GetFieldName()] = nil
@@ -1667,6 +1689,18 @@ func convertDefaultValue(value interface{}, dataType schemapb.DataType) (*schema
}
return data, nil
case schemapb.DataType_Timestamptz:
v, ok := value.(float64)
if !ok {
return nil, merr.WrapErrParameterInvalid("string", value, "Wrong defaultValue type")
}
data := &schemapb.ValueField{
Data: &schemapb.ValueField_TimestamptzData{
TimestamptzData: int64(v),
},
}
return data, nil
case schemapb.DataType_String, schemapb.DataType_VarChar:
v, ok := value.(string)
if !ok {
internal/distributed/proxy/httpserver/wrap_request.go
+16
View File
@@ -215,6 +215,22 @@ func (f *FieldData) AsSchemapb() (*schemapb.FieldData, error) {
},
}
case schemapb.DataType_Timestamptz:
data := []int64{}
err := json.Unmarshal(raw, &data)
if err != nil {
return nil, newFieldDataError(f.FieldName, err)
}
ret.Field = &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_TimestamptzData{
TimestamptzData: &schemapb.TimestamptzArray{
Data: data,
},
},
},
}
case schemapb.DataType_FloatVector:
wrappedData := [][]float32{}
err := json.Unmarshal(raw, &wrappedData)
internal/proxy/validate_util.go
+29 -1
View File
@@ -131,7 +131,8 @@ func (v *validateUtil) Validate(data []*schemapb.FieldData, helper *typeutil.Sch
case schemapb.DataType_ArrayOfStruct:
panic("unreachable, array of struct should have been flattened")
case schemapb.DataType_Timestamptz:
// TODO: Add check logic for timestamptz data
default:
}
}
@@ -376,6 +377,12 @@ func (v *validateUtil) fillWithNullValue(field *schemapb.FieldData, fieldSchema
return err
}
case *schemapb.ScalarField_TimestamptzData:
sd.TimestamptzData.Data, err = fillWithNullValueImpl(sd.TimestamptzData.Data, field.GetValidData())
if err != nil {
return err
}
case *schemapb.ScalarField_StringData:
sd.StringData.Data, err = fillWithNullValueImpl(sd.StringData.Data, field.GetValidData())
if err != nil {
@@ -466,6 +473,17 @@ func (v *validateUtil) fillWithDefaultValue(field *schemapb.FieldData, fieldSche
return err
}
case *schemapb.ScalarField_TimestamptzData:
if len(field.GetValidData()) != numRows {
msg := fmt.Sprintf("the length of valid_data of field(%s) is wrong", field.GetFieldName())
return merr.WrapErrParameterInvalid(numRows, len(field.GetValidData()), msg)
}
defaultValue := fieldSchema.GetDefaultValue().GetTimestamptzData()
sd.TimestamptzData.Data, err = fillWithDefaultValueImpl(sd.TimestamptzData.Data, defaultValue, field.GetValidData())
if err != nil {
return nil
}
case *schemapb.ScalarField_StringData:
if len(field.GetValidData()) != numRows {
msg := fmt.Sprintf("the length of valid_data of field(%s) is wrong", field.GetFieldName())
@@ -775,6 +793,16 @@ func (v *validateUtil) checkDoubleFieldData(field *schemapb.FieldData, fieldSche
return nil
}
func (v *validateUtil) checkTimestamptzFieldData(field *schemapb.FieldData, fieldSchema *schemapb.FieldSchema) error {
data := field.GetScalars().GetTimestamptzData().GetData()
if data == nil && fieldSchema.GetDefaultValue() == nil && !fieldSchema.GetNullable() {
msg := fmt.Sprintf("field '%v' is illegal, array type mismatch", field.GetFieldName())
return merr.WrapErrParameterInvalid("need long int array", "got nil", msg)
}
// TODO: Additional checks?
return nil
}
func (v *validateUtil) checkArrayElement(array *schemapb.ArrayArray, field *schemapb.FieldSchema) error {
switch field.GetElementType() {
case schemapb.DataType_Bool:
internal/querynodev2/delegator/segment_pruner.go
+1 -1
View File
@@ -319,7 +319,7 @@ func FilterSegmentsOnScalarField(partitionStats *storage.PartitionStatsSnapshot,
targetRange := tRange.ToIntRange()
statRange := exprutil.NewIntRange(min.GetValue().(int64), max.GetValue().(int64), true, true)
return exprutil.IntRangeOverlap(targetRange, statRange)
// todo: add float/double pruner
// todo: add float/double/timestmaptz pruner
case schemapb.DataType_String, schemapb.DataType_VarChar:
targetRange := tRange.ToStrRange()
statRange := exprutil.NewStrRange(min.GetValue().(string), max.GetValue().(string), true, true)
internal/querynodev2/segments/result_sorter.go
+3
View File
@@ -58,6 +58,9 @@ func swapFieldData(field *schemapb.FieldData, i int, j int) {
case *schemapb.ScalarField_DoubleData:
data := sd.DoubleData.Data
data[i], data[j] = data[j], data[i]
case *schemapb.ScalarField_TimestamptzData:
data := sd.TimestamptzData.Data
data[i], data[j] = data[j], data[i]
case *schemapb.ScalarField_StringData:
data := sd.StringData.Data
data[i], data[j] = data[j], data[i]
internal/querynodev2/segments/utils.go
+1 -1
View File
@@ -76,7 +76,7 @@ func getPKsFromRowBasedInsertMsg(msg *msgstream.InsertMsg, schema *schemapb.Coll
offset += 2
case schemapb.DataType_Int32:
offset += 4
case schemapb.DataType_Int64:
case schemapb.DataType_Timestamptz, schemapb.DataType_Int64:
offset += 8
case schemapb.DataType_Float:
offset += 4
internal/querynodev2/services.go
+1 -1
View File
@@ -975,7 +975,7 @@ func (node *QueryNode) Query(ctx context.Context, req *querypb.QueryRequest) (*i
log.Debug("received query request",
zap.Int64s("outputFields", req.GetReq().GetOutputFieldsId()),
zap.Int64s("segmentIDs", req.GetSegmentIDs()),
zap.Int64s("segmentIDs", req.GetSegmentIDs()), // should be empty
zap.Uint64("guaranteeTimestamp", req.GetReq().GetGuaranteeTimestamp()),
zap.Uint64("mvccTimestamp", req.GetReq().GetMvccTimestamp()),
zap.Bool("isCount", req.GetReq().GetIsCount()),
internal/rootcoord/util.go
+4
View File
@@ -446,6 +446,10 @@ func checkFieldSchema(fieldSchemas []*schemapb.FieldSchema) error {
if dtype != schemapb.DataType_Double {
return errTypeMismatch(fieldSchema.GetName(), dtype.String(), "DataType_Double")
}
case *schemapb.ValueField_TimestamptzData:
if dtype != schemapb.DataType_Timestamptz {
return errTypeMismatch(fieldSchema.GetName(), dtype.String(), "DataType_Timestamptz")
}
case *schemapb.ValueField_StringData:
if dtype != schemapb.DataType_VarChar {
return errTypeMismatch(fieldSchema.GetName(), dtype.String(), "DataType_VarChar")
internal/storage/arrow_util.go
+7
View File
@@ -285,6 +285,13 @@ func GenerateEmptyArrayFromSchema(schema *schemapb.FieldSchema, numRows int) (ar
bd.AppendValues(
lo.RepeatBy(numRows, func(_ int) float64 { return schema.GetDefaultValue().GetDoubleData() }),
nil)
case schemapb.DataType_Timestamptz:
bd := builder.(*array.Int64Builder)
bd.AppendValues(
lo.RepeatBy(numRows, func(_ int) int64 { return schema.GetDefaultValue().GetTimestamptzData() }),
nil)
case schemapb.DataType_VarChar, schemapb.DataType_String:
bd := builder.(*array.StringBuilder)
bd.AppendValues(
internal/storage/data_codec.go
+16
View File
@@ -388,6 +388,10 @@ func AddFieldDataToPayload(eventWriter *insertEventWriter, dataType schemapb.Dat
if err = eventWriter.AddDoubleToPayload(singleData.(*DoubleFieldData).Data, singleData.(*DoubleFieldData).ValidData); err != nil {
return err
}
case schemapb.DataType_Timestamptz:
if err = eventWriter.AddTimestamptzToPayload(singleData.(*TimestamptzFieldData).Data, singleData.(*TimestamptzFieldData).ValidData); err != nil {
return err
}
case schemapb.DataType_String, schemapb.DataType_VarChar, schemapb.DataType_Text:
for i, singleString := range singleData.(*StringFieldData).Data {
isValid := true
@@ -647,6 +651,18 @@ func AddInsertData(dataType schemapb.DataType, data interface{}, insertData *Ins
insertData.Data[fieldID] = doubleFieldData
return len(singleData), nil
case schemapb.DataType_Timestamptz:
singleData := data.([]int64)
if fieldData == nil {
fieldData = &TimestamptzFieldData{Data: make([]int64, 0, rowNum)}
}
timestamptzFieldData := fieldData.(*TimestamptzFieldData)
timestamptzFieldData.Data = append(timestamptzFieldData.Data, singleData...)
timestamptzFieldData.ValidData = append(timestamptzFieldData.ValidData, validData...)
insertData.Data[fieldID] = timestamptzFieldData
return len(singleData), nil
case schemapb.DataType_String, schemapb.DataType_VarChar, schemapb.DataType_Text:
singleData := data.([]string)
if fieldData == nil {
internal/storage/data_sorter.go
+3
View File
@@ -82,6 +82,9 @@ func (ds *DataSorter) Swap(i, j int) {
case schemapb.DataType_Double:
data := singleData.(*DoubleFieldData).Data
data[i], data[j] = data[j], data[i]
case schemapb.DataType_Timestamptz:
data := singleData.(*TimestamptzFieldData).Data
data[i], data[j] = data[j], data[i]
case schemapb.DataType_String, schemapb.DataType_VarChar:
data := singleData.(*StringFieldData).Data
data[i], data[j] = data[j], data[i]
internal/storage/field_stats.go
+15 -2
View File
@@ -100,7 +100,7 @@ func (stats *FieldStats) UnmarshalJSON(data []byte) error {
stats.Max = &Int32FieldValue{}
stats.Min = &Int32FieldValue{}
isScalarField = true
case schemapb.DataType_Int64:
case schemapb.DataType_Int64, schemapb.DataType_Timestamptz:
stats.Max = &Int64FieldValue{}
stats.Min = &Int64FieldValue{}
isScalarField = true
@@ -287,6 +287,19 @@ func (stats *FieldStats) UpdateByMsgs(msgs FieldData) {
common.Endian.PutUint64(b, uint64(doubleValue))
stats.BF.Add(b)
}
case schemapb.DataType_Timestamptz:
data := msgs.(*TimestamptzFieldData).Data
// return error: msgs must has one element at least
if len(data) < 1 {
return
}
b := make([]byte, 8)
for _, int64Value := range data {
pk := NewInt64FieldValue(int64Value)
stats.UpdateMinMax(pk)
common.Endian.PutUint64(b, uint64(int64Value))
stats.BF.Add(b)
}
case schemapb.DataType_String:
data := msgs.(*StringFieldData).Data
// return error: msgs must has one element at least
@@ -332,7 +345,7 @@ func (stats *FieldStats) Update(pk ScalarFieldValue) {
b := make([]byte, 8)
common.Endian.PutUint64(b, uint64(data))
stats.BF.Add(b)
case schemapb.DataType_Int64:
case schemapb.DataType_Int64, schemapb.DataType_Timestamptz:
data := pk.GetValue().(int64)
b := make([]byte, 8)
common.Endian.PutUint64(b, uint64(data))
internal/storage/field_value.go
+5
View File
@@ -1064,6 +1064,9 @@ func NewScalarFieldValueFromGenericValue(dtype schemapb.DataType, gVal *planpb.G
case schemapb.DataType_Double:
floatVal := gVal.Val.(*planpb.GenericValue_FloatVal)
return NewDoubleFieldValue(floatVal.FloatVal), nil
case schemapb.DataType_Timestamptz:
i64VAl := gVal.Val.(*planpb.GenericValue_Int64Val)
return NewInt64FieldValue(i64VAl.Int64Val), nil
case schemapb.DataType_String:
strVal := gVal.Val.(*planpb.GenericValue_StringVal)
return NewStringFieldValue(strVal.StringVal), nil
@@ -1090,6 +1093,8 @@ func NewScalarFieldValue(dtype schemapb.DataType, data interface{}) ScalarFieldV
return NewFloatFieldValue(data.(float32))
case schemapb.DataType_Double:
return NewDoubleFieldValue(data.(float64))
case schemapb.DataType_Timestamptz:
return NewInt64FieldValue(data.(int64))
case schemapb.DataType_String:
return NewStringFieldValue(data.(string))
case schemapb.DataType_VarChar:
internal/storage/insert_data.go
+82
View File
@@ -327,6 +327,15 @@ func NewFieldData(dataType schemapb.DataType, fieldSchema *schemapb.FieldSchema,
data.ValidData = make([]bool, 0, cap)
}
return data, nil
case schemapb.DataType_Timestamptz:
data := &TimestamptzFieldData{
Data: make([]int64, 0, cap),
Nullable: fieldSchema.GetNullable(),
}
if fieldSchema.GetNullable() {
data.ValidData = make([]bool, 0, cap)
}
return data, nil
case schemapb.DataType_JSON:
data := &JSONFieldData{
Data: make([][]byte, 0, cap),
@@ -419,6 +428,11 @@ type JSONFieldData struct {
ValidData []bool
Nullable bool
}
type TimestamptzFieldData struct {
Data []int64
ValidData []bool
Nullable bool
}
type BinaryVectorFieldData struct {
Data []byte
Dim int
@@ -469,6 +483,7 @@ func (data *Int32FieldData) RowNum() int { return len(data.Data) }
func (data *Int64FieldData) RowNum() int { return len(data.Data) }
func (data *FloatFieldData) RowNum() int { return len(data.Data) }
func (data *DoubleFieldData) RowNum() int { return len(data.Data) }
func (data *TimestamptzFieldData) RowNum() int { return len(data.Data) }
func (data *StringFieldData) RowNum() int { return len(data.Data) }
func (data *ArrayFieldData) RowNum() int { return len(data.Data) }
func (data *JSONFieldData) RowNum() int { return len(data.Data) }
@@ -520,6 +535,13 @@ func (data *Int64FieldData) GetRow(i int) any {
return data.Data[i]
}
func (data *TimestamptzFieldData) GetRow(i int) any {
if data.GetNullable() && !data.ValidData[i] {
return nil
}
return data.Data[i]
}
func (data *FloatFieldData) GetRow(i int) any {
if data.GetNullable() && !data.ValidData[i] {
return nil
@@ -590,6 +612,7 @@ func (data *Int32FieldData) GetDataRows() any { return data.Data }
func (data *Int64FieldData) GetDataRows() any { return data.Data }
func (data *FloatFieldData) GetDataRows() any { return data.Data }
func (data *DoubleFieldData) GetDataRows() any { return data.Data }
func (data *TimestamptzFieldData) GetDataRows() any { return data.Data }
func (data *StringFieldData) GetDataRows() any { return data.Data }
func (data *ArrayFieldData) GetDataRows() any { return data.Data }
func (data *JSONFieldData) GetDataRows() any { return data.Data }
@@ -721,6 +744,23 @@ func (data *DoubleFieldData) AppendRow(row interface{}) error {
return nil
}
func (data *TimestamptzFieldData) AppendRow(row interface{}) error {
if data.GetNullable() && row == nil {
data.Data = append(data.Data, make([]int64, 1)...)
data.ValidData = append(data.ValidData, false)
return nil
}
v, ok := row.(int64)
if !ok {
return merr.WrapErrParameterInvalid("timestamptz", row, "Wrong row type")
}
if data.GetNullable() {
data.ValidData = append(data.ValidData, true)
}
data.Data = append(data.Data, v)
return nil
}
func (data *StringFieldData) AppendRow(row interface{}) error {
if data.GetNullable() && row == nil {
data.Data = append(data.Data, make([]string, 1)...)
@@ -898,6 +938,14 @@ func (data *DoubleFieldData) AppendRows(dataRows interface{}, validDataRows inte
return data.AppendValidDataRows(validDataRows)
}
func (data *TimestamptzFieldData) AppendRows(dataRows interface{}, validDataRows interface{}) error {
err := data.AppendDataRows(dataRows)
if err != nil {
return err
}
return data.AppendValidDataRows(validDataRows)
}
func (data *StringFieldData) AppendRows(dataRows interface{}, validDataRows interface{}) error {
err := data.AppendDataRows(dataRows)
if err != nil {
@@ -1045,6 +1093,15 @@ func (data *DoubleFieldData) AppendDataRows(rows interface{}) error {
return nil
}
func (data *TimestamptzFieldData) AppendDataRows(rows interface{}) error {
v, ok := rows.([]int64)
if !ok {
return merr.WrapErrParameterInvalid("[]timestamptz", rows, "Wrong rows type")
}
data.Data = append(data.Data, v...)
return nil
}
func (data *StringFieldData) AppendDataRows(rows interface{}) error {
v, ok := rows.([]string)
if !ok {
@@ -1241,6 +1298,18 @@ func (data *DoubleFieldData) AppendValidDataRows(rows interface{}) error {
return nil
}
func (data *TimestamptzFieldData) AppendValidDataRows(rows interface{}) error {
if rows == nil {
return nil
}
v, ok := rows.([]bool)
if !ok {
return merr.WrapErrParameterInvalid("[]bool", rows, "Wrong rows type")
}
data.ValidData = append(data.ValidData, v...)
return nil
}
func (data *StringFieldData) AppendValidDataRows(rows interface{}) error {
if rows == nil {
return nil
@@ -1400,6 +1469,11 @@ func (data *FloatFieldData) GetMemorySize() int {
func (data *DoubleFieldData) GetMemorySize() int {
return binary.Size(data.Data) + binary.Size(data.ValidData) + binary.Size(data.Nullable)
}
func (data *TimestamptzFieldData) GetMemorySize() int {
return binary.Size(data.Data) + binary.Size(data.ValidData) + binary.Size(data.Nullable)
}
func (data *BinaryVectorFieldData) GetMemorySize() int { return binary.Size(data.Data) + 4 }
func (data *FloatVectorFieldData) GetMemorySize() int { return binary.Size(data.Data) + 4 }
func (data *Float16VectorFieldData) GetMemorySize() int { return binary.Size(data.Data) + 4 }
@@ -1449,6 +1523,9 @@ func (data *Int32FieldData) GetDataType() schemapb.DataType { return schemapb.D
func (data *Int64FieldData) GetDataType() schemapb.DataType { return schemapb.DataType_Int64 }
func (data *FloatFieldData) GetDataType() schemapb.DataType { return schemapb.DataType_Float }
func (data *DoubleFieldData) GetDataType() schemapb.DataType { return schemapb.DataType_Double }
func (data *TimestamptzFieldData) GetDataType() schemapb.DataType {
return schemapb.DataType_Timestamptz
}
func (data *StringFieldData) GetDataType() schemapb.DataType { return data.DataType }
func (data *ArrayFieldData) GetDataType() schemapb.DataType { return schemapb.DataType_Array }
func (data *JSONFieldData) GetDataType() schemapb.DataType { return schemapb.DataType_JSON }
@@ -1532,6 +1609,7 @@ func (data *Int32FieldData) GetRowSize(i int) int { return 4 }
func (data *Int64FieldData) GetRowSize(i int) int { return 8 }
func (data *FloatFieldData) GetRowSize(i int) int { return 4 }
func (data *DoubleFieldData) GetRowSize(i int) int { return 8 }
func (data *TimestamptzFieldData) GetRowSize(i int) int { return 8 }
func (data *BinaryVectorFieldData) GetRowSize(i int) int { return data.Dim / 8 }
func (data *FloatVectorFieldData) GetRowSize(i int) int { return data.Dim * 4 }
func (data *Float16VectorFieldData) GetRowSize(i int) int { return data.Dim * 2 }
@@ -1597,6 +1675,10 @@ func (data *DoubleFieldData) GetNullable() bool {
return data.Nullable
}
func (data *TimestamptzFieldData) GetNullable() bool {
return data.Nullable
}
func (data *BFloat16VectorFieldData) GetNullable() bool {
return false
}
internal/storage/payload.go
+2
View File
@@ -35,6 +35,7 @@ type PayloadWriterInterface interface {
AddInt64ToPayload([]int64, []bool) error
AddFloatToPayload([]float32, []bool) error
AddDoubleToPayload([]float64, []bool) error
AddTimestamptzToPayload([]int64, []bool) error
AddOneStringToPayload(string, bool) error
AddOneArrayToPayload(*schemapb.ScalarField, bool) error
AddOneJSONToPayload([]byte, bool) error
@@ -64,6 +65,7 @@ type PayloadReaderInterface interface {
GetInt64FromPayload() ([]int64, []bool, error)
GetFloatFromPayload() ([]float32, []bool, error)
GetDoubleFromPayload() ([]float64, []bool, error)
GetTimestamptzFromPayload() ([]int64, []bool, error)
GetStringFromPayload() ([]string, []bool, error)
GetArrayFromPayload() ([]*schemapb.ScalarField, []bool, error)
GetVectorArrayFromPayload() ([]*schemapb.VectorField, error)
internal/storage/payload_reader.go
+34
View File
@@ -74,6 +74,9 @@ func (r *PayloadReader) GetDataFromPayload() (interface{}, []bool, int, error) {
case schemapb.DataType_Double:
val, validData, err := r.GetDoubleFromPayload()
return val, validData, 0, err
case schemapb.DataType_Timestamptz:
val, validData, err := r.GetTimestamptzFromPayload()
return val, validData, 0, err
case schemapb.DataType_BinaryVector:
val, dim, err := r.GetBinaryVectorFromPayload()
return val, nil, dim, err
@@ -369,6 +372,37 @@ func (r *PayloadReader) GetDoubleFromPayload() ([]float64, []bool, error) {
return values, nil, nil
}
func (r *PayloadReader) GetTimestamptzFromPayload() ([]int64, []bool, error) {
if r.colType != schemapb.DataType_Timestamptz {
return nil, nil, merr.WrapErrParameterInvalidMsg(fmt.Sprintf("failed to get timestamptz from datatype %v", r.colType.String()))
}
values := make([]int64, r.numRows)
if r.nullable {
validData := make([]bool, r.numRows)
valuesRead, err := ReadData[int64, *array.Int64](r.reader, values, validData, r.numRows)
if err != nil {
return nil, nil, err
}
if valuesRead != r.numRows {
return nil, nil, merr.WrapErrParameterInvalid(r.numRows, valuesRead, "valuesRead is not equal to rows")
}
return values, validData, nil
}
valuesRead, err := ReadDataFromAllRowGroups[int64, *file.Int64ColumnChunkReader](r.reader, values, 0, r.numRows)
if err != nil {
return nil, nil, err
}
if valuesRead != r.numRows {
return nil, nil, merr.WrapErrParameterInvalid(r.numRows, valuesRead, "valuesRead is not equal to rows")
}
return values, nil, nil
}
func (r *PayloadReader) GetStringFromPayload() ([]string, []bool, error) {
if r.colType != schemapb.DataType_String && r.colType != schemapb.DataType_VarChar {
return nil, nil, merr.WrapErrParameterInvalidMsg(fmt.Sprintf("failed to get string from datatype %v", r.colType.String()))
internal/storage/payload_writer.go
+35 -1
View File
@@ -150,6 +150,12 @@ func (w *NativePayloadWriter) AddDataToPayload(data interface{}, validData []boo
return merr.WrapErrParameterInvalidMsg("incorrect data type")
}
return w.AddDoubleToPayload(val, validData)
case schemapb.DataType_Timestamptz:
val, ok := data.([]int64)
if !ok {
return merr.WrapErrParameterInvalidMsg("incorrect data type")
}
return w.AddTimestamptzToPayload(val, validData)
case schemapb.DataType_String, schemapb.DataType_VarChar:
val, ok := data.(string)
if !ok {
@@ -485,6 +491,34 @@ func (w *NativePayloadWriter) AddDoubleToPayload(data []float64, validData []boo
return nil
}
func (w *NativePayloadWriter) AddTimestamptzToPayload(data []int64, validData []bool) error {
if w.finished {
return errors.New("can't append data to finished int64 payload")
}
if len(data) == 0 {
return errors.New("can't add empty msgs into int64 payload")
}
if !w.nullable && len(validData) != 0 {
msg := fmt.Sprintf("length of validData(%d) must be 0 when not nullable", len(validData))
return merr.WrapErrParameterInvalidMsg(msg)
}
if w.nullable && len(data) != len(validData) {
msg := fmt.Sprintf("length of validData(%d) must equal to data(%d) when nullable", len(validData), len(data))
return merr.WrapErrParameterInvalidMsg(msg)
}
builder, ok := w.builder.(*array.Int64Builder)
if !ok {
return errors.New("failed to cast Int64Builder")
}
builder.AppendValues(data, validData)
return nil
}
func (w *NativePayloadWriter) AddOneStringToPayload(data string, isValid bool) error {
if w.finished {
return errors.New("can't append data to finished string payload")
@@ -800,7 +834,7 @@ func MilvusDataTypeToArrowType(dataType schemapb.DataType, dim int) arrow.DataTy
return &arrow.Int16Type{}
case schemapb.DataType_Int32:
return &arrow.Int32Type{}
case schemapb.DataType_Int64:
case schemapb.DataType_Int64, schemapb.DataType_Timestamptz:
return &arrow.Int64Type{}
case schemapb.DataType_Float:
return &arrow.Float32Type{}
internal/storage/print_binlog.go
+8
View File
@@ -279,6 +279,14 @@ func printPayloadValues(colType schemapb.DataType, reader PayloadReaderInterface
for i, v := range val {
fmt.Printf("\t\t%d : %v\n", i, v)
}
case schemapb.DataType_Timestamptz:
val, _, err := reader.GetTimestamptzFromPayload()
if err != nil {
return err
}
for i, v := range val {
fmt.Printf("\t\t%d : %d\n", i, v)
}
case schemapb.DataType_String, schemapb.DataType_VarChar:
rows, err := reader.GetPayloadLengthFromReader()
if err != nil {
internal/storage/serde.go
+28 -1
View File
@@ -294,6 +294,33 @@ var serdeMap = func() map[schemapb.DataType]serdeEntry {
return false
},
}
m[schemapb.DataType_Timestamptz] = serdeEntry{
arrowType: func(i int) arrow.DataType {
return arrow.PrimitiveTypes.Int64
},
deserialize: func(a arrow.Array, i int, shouldCopy bool) (any, bool) {
if a.IsNull(i) {
return nil, true
}
if arr, ok := a.(*array.Int64); ok && i < arr.Len() {
return arr.Value(i), true
}
return nil, false
},
serialize: func(b array.Builder, v any) bool {
if v == nil {
b.AppendNull()
return true
}
if builder, ok := b.(*array.Int64Builder); ok {
if v, ok := v.(int64); ok {
builder.Append(v)
return true
}
}
return false
},
}
stringEntry := serdeEntry{
arrowType: func(i int) arrow.DataType {
return arrow.BinaryTypes.String
@@ -735,7 +762,7 @@ func newSingleFieldRecordWriter(field *schemapb.FieldSchema, writer io.Writer, o
return 2
case schemapb.DataType_Int32:
return 4
case schemapb.DataType_Int64:
case schemapb.DataType_Int64, schemapb.DataType_Timestamptz:
return 8
}
}
internal/storage/utils.go
+57
View File
@@ -371,6 +371,16 @@ func readDoubleArray(blobReaders []io.Reader) []float64 {
return ret
}
func readTimestamptzArray(blobReaders []io.Reader) []int64 {
ret := make([]int64, 0)
for _, r := range blobReaders {
var v int64
ReadBinary(r, &v, schemapb.DataType_Timestamptz) // DataType is only used for logging
ret = append(ret, v)
}
return ret
}
func RowBasedInsertMsgToInsertData(msg *msgstream.InsertMsg, collSchema *schemapb.CollectionSchema, skipFunction bool) (idata *InsertData, err error) {
blobReaders := make([]io.Reader, 0)
for _, blob := range msg.RowData {
@@ -507,6 +517,11 @@ func RowBasedInsertMsgToInsertData(msg *msgstream.InsertMsg, collSchema *schemap
idata.Data[field.FieldID] = &DoubleFieldData{
Data: readDoubleArray(blobReaders),
}
case schemapb.DataType_Timestamptz:
idata.Data[field.FieldID] = &TimestamptzFieldData{
Data: readTimestamptzArray(blobReaders),
}
}
}
@@ -697,6 +712,16 @@ func ColumnBasedInsertMsgToInsertData(msg *msgstream.InsertMsg, collSchema *sche
Nullable: field.GetNullable(),
}
case schemapb.DataType_Timestamptz:
srcData := srcField.GetScalars().GetTimestamptzData().GetData()
validData := srcField.GetValidData()
fieldData = &TimestamptzFieldData{
Data: srcData,
ValidData: validData,
Nullable: field.GetNullable(),
}
case schemapb.DataType_String, schemapb.DataType_VarChar, schemapb.DataType_Text:
srcData := srcField.GetScalars().GetStringData().GetData()
validData := srcField.GetValidData()
@@ -891,6 +916,19 @@ func mergeDoubleField(data *InsertData, fid FieldID, field *DoubleFieldData) {
fieldData.ValidData = append(fieldData.ValidData, field.ValidData...)
}
func mergeTimestamptzField(data *InsertData, fid FieldID, field *TimestamptzFieldData) {
if _, ok := data.Data[fid]; !ok {
fieldData := &TimestamptzFieldData{
Data: nil,
ValidData: nil,
}
data.Data[fid] = fieldData
}
fieldData := data.Data[fid].(*TimestamptzFieldData)
fieldData.Data = append(fieldData.Data, field.Data...)
fieldData.ValidData = append(fieldData.ValidData, field.ValidData...)
}
func mergeStringField(data *InsertData, fid FieldID, field *StringFieldData) {
if _, ok := data.Data[fid]; !ok {
fieldData := &StringFieldData{
@@ -1019,6 +1057,8 @@ func MergeFieldData(data *InsertData, fid FieldID, field FieldData) {
mergeFloatField(data, fid, field)
case *DoubleFieldData:
mergeDoubleField(data, fid, field)
case *TimestamptzFieldData:
mergeTimestamptzField(data, fid, field)
case *StringFieldData:
mergeStringField(data, fid, field)
case *ArrayFieldData:
@@ -1261,6 +1301,21 @@ func TransferInsertDataToInsertRecord(insertData *InsertData) (*segcorepb.Insert
},
ValidData: rawData.ValidData,
}
case *TimestamptzFieldData:
fieldData = &schemapb.FieldData{
Type: schemapb.DataType_Timestamptz,
FieldId: fieldID,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_TimestamptzData{
TimestamptzData: &schemapb.TimestamptzArray{
Data: rawData.Data,
},
},
},
},
ValidData: rawData.ValidData,
}
case *StringFieldData:
fieldData = &schemapb.FieldData{
Type: schemapb.DataType_VarChar,
@@ -1522,6 +1577,8 @@ func GetDefaultValue(fieldSchema *schemapb.FieldSchema) interface{} {
return fieldSchema.GetDefaultValue().GetDoubleData()
case schemapb.DataType_VarChar, schemapb.DataType_String:
return fieldSchema.GetDefaultValue().GetStringData()
case schemapb.DataType_Timestamptz:
return fieldSchema.GetDefaultValue().GetTimestamptzData()
default:
// won't happen
panic(fmt.Sprintf("undefined data type:%s", fieldSchema.DataType.String()))
internal/util/function/rerank/decay_function.go
+2 -2
View File
@@ -69,7 +69,7 @@ func newDecayFunction(collSchema *schemapb.CollectionSchema, funcSchema *schemap
switch inputType {
case schemapb.DataType_Int8, schemapb.DataType_Int16, schemapb.DataType_Int32:
return newFunction[int64, int32](base, funcSchema)
case schemapb.DataType_Int64:
case schemapb.DataType_Int64, schemapb.DataType_Timestamptz:
return newFunction[int64, int64](base, funcSchema)
case schemapb.DataType_Float:
return newFunction[int64, float32](base, funcSchema)
@@ -82,7 +82,7 @@ func newDecayFunction(collSchema *schemapb.CollectionSchema, funcSchema *schemap
switch inputType {
case schemapb.DataType_Int8, schemapb.DataType_Int16, schemapb.DataType_Int32:
return newFunction[string, int32](base, funcSchema)
case schemapb.DataType_Int64:
case schemapb.DataType_Int64, schemapb.DataType_Timestamptz:
return newFunction[string, int64](base, funcSchema)
case schemapb.DataType_Float:
return newFunction[string, float32](base, funcSchema)
internal/util/function/rerank/util.go
+5
View File
@@ -344,6 +344,11 @@ func getField(inputField *schemapb.FieldData, start int64, size int64) (any, err
return inputField.GetScalars().GetDoubleData().Data[start : start+size], nil
}
return []float64{}, nil
case schemapb.DataType_Timestamptz:
if inputField.GetScalars() != nil && inputField.GetScalars().GetTimestamptzData() != nil {
return inputField.GetScalars().GetTimestamptzData().Data[start : start+size], nil
}
return []int64{}, nil
case schemapb.DataType_Bool:
if inputField.GetScalars() != nil && inputField.GetScalars().GetBoolData() != nil {
return inputField.GetScalars().GetBoolData().Data[start : start+size], nil
internal/util/importutilv2/csv/row_parser.go
+21 -1
View File
@@ -254,7 +254,7 @@ func (r *rowParser) parseEntity(field *schemapb.FieldSchema, obj string) (any, e
return 0, r.wrapTypeError(obj, field)
}
return int32(num), nil
case schemapb.DataType_Int64:
case schemapb.DataType_Int64, schemapb.DataType_Timestamptz:
num, err := strconv.ParseInt(obj, 10, 64)
if err != nil {
return 0, r.wrapTypeError(obj, field)
@@ -493,6 +493,26 @@ func (r *rowParser) arrayToFieldData(arr []interface{}, field *schemapb.FieldSch
},
},
}, nil
case schemapb.DataType_Timestamptz:
values := make([]int64, len(arr))
for i, v := range arr {
value, ok := v.(json.Number)
if !ok {
return nil, r.wrapArrayValueTypeError(arr, eleType)
}
num, err := strconv.ParseInt(value.String(), 10, 64)
if err != nil {
return nil, fmt.Errorf("failed to parse timesamptz: %w", err)
}
values[i] = num
}
return &schemapb.ScalarField{
Data: &schemapb.ScalarField_TimestamptzData{
TimestamptzData: &schemapb.TimestamptzArray{
Data: values,
},
},
}, nil
case schemapb.DataType_VarChar, schemapb.DataType_String:
values := make([]string, len(arr))
for i, v := range arr {
internal/util/importutilv2/json/row_parser.go
+21 -1
View File
@@ -277,7 +277,7 @@ func (r *rowParser) parseEntity(fieldID int64, obj any) (any, error) {
return nil, err
}
return int32(num), nil
case schemapb.DataType_Int64:
case schemapb.DataType_Int64, schemapb.DataType_Timestamptz:
value, ok := obj.(json.Number)
if !ok {
return nil, r.wrapTypeError(obj, fieldID)
@@ -586,6 +586,26 @@ func (r *rowParser) arrayToFieldData(arr []interface{}, field *schemapb.FieldSch
},
},
}, nil
case schemapb.DataType_Timestamptz:
values := make([]int64, len(arr))
for i, v := range arr {
value, ok := v.(json.Number)
if !ok {
return nil, r.wrapArrayValueTypeError(arr, eleType)
}
num, err := strconv.ParseInt(value.String(), 0, 64)
if err != nil {
return nil, fmt.Errorf("failed to parse int64: %w", err)
}
values[i] = num
}
return &schemapb.ScalarField{
Data: &schemapb.ScalarField_TimestamptzData{
TimestamptzData: &schemapb.TimestamptzArray{
Data: values,
},
},
}, nil
case schemapb.DataType_VarChar, schemapb.DataType_String:
values := make([]string, len(arr))
for i, v := range arr {
internal/util/importutilv2/numpy/field_reader.go
+5
View File
@@ -178,6 +178,11 @@ func (c *FieldReader) Next(count int64) (any, any, error) {
return nil, nil, err
}
c.readPosition += int(readCount)
case schemapb.DataType_Timestamptz:
data, err = ReadN[int64](c.reader, c.order, readCount)
if err != nil {
return nil, nil, err
}
case schemapb.DataType_VarChar:
data, err = c.ReadString(readCount)
c.readPosition += int(readCount)
internal/util/importutilv2/numpy/util.go
+1 -1
View File
@@ -240,7 +240,7 @@ func convertNumpyType(typeStr string) (schemapb.DataType, error) {
return schemapb.DataType_Int16, nil
case "i4", "<i4", "|i4", ">i4", "int32":
return schemapb.DataType_Int32, nil
case "i8", "<i8", "|i8", ">i8", "int64":
case "i8", "<i8", "|i8", ">i8", "int64": // todo: how to handle Timestamptz?
return schemapb.DataType_Int64, nil
case "f4", "<f4", "|f4", ">f4", "float32":
return schemapb.DataType_Float, nil
internal/util/importutilv2/parquet/field_reader.go
+42 -1
View File
@@ -103,7 +103,7 @@ func (c *FieldReader) Next(count int64) (any, any, error) {
}
data, err := ReadIntegerOrFloatData[int32](c, count)
return data, nil, err
case schemapb.DataType_Int64:
case schemapb.DataType_Int64, schemapb.DataType_Timestamptz:
if c.field.GetNullable() || c.field.GetDefaultValue() != nil {
return ReadNullableIntegerOrFloatData[int64](c, count)
}
@@ -1519,6 +1519,26 @@ func ReadArrayData(pcr *FieldReader, count int64) (any, error) {
},
})
}
case schemapb.DataType_Timestamptz:
int64Array, err := ReadIntegerOrFloatArrayData[int64](pcr, count)
if err != nil {
return nil, err
}
if int64Array == nil {
return nil, nil
}
for _, elementArray := range int64Array.([][]int64) {
if err = common.CheckArrayCapacity(len(elementArray), maxCapacity, pcr.field); err != nil {
return nil, err
}
data = append(data, &schemapb.ScalarField{
Data: &schemapb.ScalarField_TimestamptzData{
TimestamptzData: &schemapb.TimestamptzArray{
Data: elementArray,
},
},
})
}
case schemapb.DataType_VarChar, schemapb.DataType_String:
stringArray, err := ReadStringArrayData(pcr, count)
if err != nil {
@@ -1701,6 +1721,27 @@ func ReadNullableArrayData(pcr *FieldReader, count int64) (any, []bool, error) {
})
}
return data, validData, nil
case schemapb.DataType_Timestamptz:
int64Array, validData, err := ReadNullableIntegerOrFloatArrayData[int64](pcr, count)
if err != nil {
return nil, nil, err
}
if int64Array == nil {
return nil, nil, nil
}
for _, elementArray := range int64Array.([][]int64) {
if err = common.CheckArrayCapacity(len(elementArray), maxCapacity, pcr.field); err != nil {
return nil, nil, err
}
data = append(data, &schemapb.ScalarField{
Data: &schemapb.ScalarField_TimestamptzData{
TimestamptzData: &schemapb.TimestamptzArray{
Data: elementArray,
},
},
})
}
return data, validData, nil
case schemapb.DataType_VarChar, schemapb.DataType_String:
stringArray, validData, err := ReadNullableStringArrayData(pcr, count)
if err != nil {
internal/util/importutilv2/parquet/util.go
+1 -1
View File
@@ -241,7 +241,7 @@ func convertToArrowDataType(field *schemapb.FieldSchema, isArray bool) (arrow.Da
return &arrow.Int16Type{}, nil
case schemapb.DataType_Int32:
return &arrow.Int32Type{}, nil
case schemapb.DataType_Int64:
case schemapb.DataType_Int64, schemapb.DataType_Timestamptz:
return &arrow.Int64Type{}, nil
case schemapb.DataType_Float:
return &arrow.Float32Type{}, nil
internal/util/nullutil/nullutil.go
+2
View File
@@ -53,6 +53,8 @@ func GetDefaultValue(field *schemapb.FieldSchema) (any, error) {
return field.GetDefaultValue().GetFloatData(), nil
case schemapb.DataType_Double:
return field.GetDefaultValue().GetDoubleData(), nil
case schemapb.DataType_Timestamptz:
return field.GetDefaultValue().GetTimestamptzData(), nil
case schemapb.DataType_String, schemapb.DataType_VarChar:
return field.GetDefaultValue().GetStringData(), nil
default:
pkg/util/funcutil/func.go
+4
View File
@@ -417,6 +417,8 @@ func GetNumRowOfFieldDataWithSchema(fieldData *schemapb.FieldData, helper *typeu
fieldNumRows = getNumRowsOfScalarField(fieldData.GetScalars().GetFloatData().GetData())
case schemapb.DataType_Double:
fieldNumRows = getNumRowsOfScalarField(fieldData.GetScalars().GetDoubleData().GetData())
case schemapb.DataType_Timestamptz:
fieldNumRows = getNumRowsOfScalarField(fieldData.GetScalars().GetTimestamptzData().GetData())
case schemapb.DataType_String, schemapb.DataType_VarChar, schemapb.DataType_Text:
fieldNumRows = getNumRowsOfScalarField(fieldData.GetScalars().GetStringData().GetData())
case schemapb.DataType_Array:
@@ -482,6 +484,8 @@ func GetNumRowOfFieldData(fieldData *schemapb.FieldData) (uint64, error) {
fieldNumRows = getNumRowsOfScalarField(scalarField.GetFloatData().Data)
case *schemapb.ScalarField_DoubleData:
fieldNumRows = getNumRowsOfScalarField(scalarField.GetDoubleData().Data)
case *schemapb.ScalarField_TimestamptzData:
fieldNumRows = getNumRowsOfScalarField(scalarField.GetTimestamptzData().Data)
case *schemapb.ScalarField_StringData:
fieldNumRows = getNumRowsOfScalarField(scalarField.GetStringData().Data)
case *schemapb.ScalarField_ArrayData:
pkg/util/typeutil/gen_empty_field_data.go
+2
View File
@@ -261,6 +261,8 @@ func GenEmptyFieldData(field *schemapb.FieldSchema) (*schemapb.FieldData, error)
return genEmptyFloatFieldData(field), nil
case schemapb.DataType_Double:
return genEmptyDoubleFieldData(field), nil
case schemapb.DataType_Timestamptz:
return genEmptyLongFieldData(field), nil
case schemapb.DataType_VarChar, schemapb.DataType_Text:
return genEmptyVarCharFieldData(field), nil
case schemapb.DataType_Array:
pkg/util/typeutil/schema.go
+18 -2
View File
@@ -112,7 +112,7 @@ func estimateSizeBy(schema *schemapb.CollectionSchema, policy getVariableFieldLe
res += 2
case schemapb.DataType_Int32, schemapb.DataType_Float:
res += 4
case schemapb.DataType_Int64, schemapb.DataType_Double:
case schemapb.DataType_Int64, schemapb.DataType_Double, schemapb.DataType_Timestamptz:
res += 8
case schemapb.DataType_VarChar, schemapb.DataType_Text, schemapb.DataType_Array, schemapb.DataType_JSON:
maxLengthPerRow, err := getVarFieldLength(fs, policy)
@@ -192,6 +192,8 @@ func CalcScalarSize(column *schemapb.FieldData) int {
res += len(column.GetScalars().GetFloatData().GetData()) * 4
case schemapb.DataType_Double:
res += len(column.GetScalars().GetDoubleData().GetData()) * 8
case schemapb.DataType_Timestamptz:
res += len(column.GetScalars().GetTimestamptzData().GetData()) * 8
case schemapb.DataType_VarChar, schemapb.DataType_Text:
for _, str := range column.GetScalars().GetStringData().GetData() {
res += len(str)
@@ -246,7 +248,7 @@ func EstimateEntitySize(fieldsData []*schemapb.FieldData, rowOffset int) (int, e
res += 2
case schemapb.DataType_Int32, schemapb.DataType_Float:
res += 4
case schemapb.DataType_Int64, schemapb.DataType_Double:
case schemapb.DataType_Int64, schemapb.DataType_Double, schemapb.DataType_Timestamptz:
res += 8
case schemapb.DataType_VarChar, schemapb.DataType_Text:
if rowOffset >= len(fs.GetScalars().GetStringData().GetData()) {
@@ -885,6 +887,18 @@ func AppendFieldData(dst, src []*schemapb.FieldData, idx int64) (appendSize int6
}
/* #nosec G103 */
appendSize += int64(unsafe.Sizeof(srcScalar.JsonData.Data[idx]))
case *schemapb.ScalarField_TimestamptzData:
if dstScalar.GetTimestamptzData() == nil {
dstScalar.Data = &schemapb.ScalarField_TimestamptzData{
TimestamptzData: &schemapb.TimestamptzArray{
Data: []int64{srcScalar.TimestamptzData.Data[idx]},
},
}
} else {
dstScalar.GetTimestamptzData().Data = append(dstScalar.GetTimestamptzData().Data, srcScalar.TimestamptzData.Data[idx])
}
/* #nosec G103 */
appendSize += int64(unsafe.Sizeof(srcScalar.TimestamptzData.Data[idx]))
default:
log.Error("Not supported field type", zap.String("field type", fieldData.Type.String()))
}
@@ -1794,6 +1808,8 @@ func getData(field *schemapb.FieldData, idx int) any {
return field.GetScalars().GetFloatData().GetData()[idx]
case schemapb.DataType_Double:
return field.GetScalars().GetDoubleData().GetData()[idx]
case schemapb.DataType_Timestamptz:
return field.GetScalars().GetTimestamptzData().GetData()[idx]
case schemapb.DataType_VarChar, schemapb.DataType_Text:
return field.GetScalars().GetStringData().GetData()[idx]
case schemapb.DataType_FloatVector: