milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fix arrow builder nullptr check in FinishPayloadWriter (#17873) 

Signed-off-by: Letian Jiang <letian.jiang@zilliz.com>
pull/17875/head
Letian Jiang 2022-06-28 20:04:17 +08:00 committed by GitHub
parent 32c70b7f86
commit 4ae1ca2cac
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 40 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
internal/core/src/storage/parquet_c.cpp
View File

@ -305,6 +305,12 @@ FinishPayloadWriter(CPayloadWriter payloadWriter) {
st.error_msg = nullptr;
auto p = reinterpret_cast<wrapper::PayloadWriter*>(payloadWriter);
if (p->builder == nullptr) {
if (p->dimension == wrapper::EMPTY_DIMENSION) {
// For FloatVector/BinaryVector datatype, the builder is lazily inited.
// Since wrapper::EMPTY_DIMENSION indicates the builder is not inited,
// we simply return success here.
return st;
}
st.error_code = static_cast<int>(ErrorCode::UNEXPECTED_ERROR);
st.error_msg = ErrorMsg("arrow builder is nullptr");
return st;

30
internal/core/unittest/test_parquet_c.cpp
View File

@ -248,6 +248,21 @@ TEST(wrapper, binary_vector) {
ReleasePayloadReader(reader);
}
TEST(wrapper, binary_vector_empty) {
auto payload = NewPayloadWriter(ColumnType::VECTOR_BINARY);
auto st = FinishPayloadWriter(payload);
ASSERT_EQ(st.error_code, ErrorCode::SUCCESS);
auto cb = GetPayloadBufferFromWriter(payload);
ASSERT_EQ(cb.length, 0);
ASSERT_EQ(cb.data, nullptr);
auto nums = GetPayloadLengthFromWriter(payload);
ASSERT_EQ(nums, 0);
auto reader = NewPayloadReader(ColumnType::VECTOR_BINARY, (uint8_t*)cb.data, cb.length);
ASSERT_EQ(reader, nullptr);
ReleasePayloadWriter(payload);
ReleasePayloadReader(reader);
}
TEST(wrapper, float_vector) {
auto payload = NewPayloadWriter(ColumnType::VECTOR_FLOAT);
float data[] = {1, 2, 3, 4, 5, 6, 7, 8};
@ -282,6 +297,21 @@ TEST(wrapper, float_vector) {
ReleasePayloadReader(reader);
}
TEST(wrapper, float_vector_empty) {
auto payload = NewPayloadWriter(ColumnType::VECTOR_FLOAT);
auto st = FinishPayloadWriter(payload);
ASSERT_EQ(st.error_code, ErrorCode::SUCCESS);
auto cb = GetPayloadBufferFromWriter(payload);
ASSERT_EQ(cb.length, 0);
ASSERT_EQ(cb.data, nullptr);
auto nums = GetPayloadLengthFromWriter(payload);
ASSERT_EQ(nums, 0);
auto reader = NewPayloadReader(ColumnType::VECTOR_FLOAT, (uint8_t*)cb.data, cb.length);
ASSERT_EQ(reader, nullptr);
ReleasePayloadWriter(payload);
ReleasePayloadReader(reader);
}
TEST(wrapper, int8_2) {
auto payload = NewPayloadWriter(ColumnType::INT8);
int8_t data[] = {-1, 1, -100, 100};

4
internal/storage/payload_cgo_test.go
View File

@ -583,7 +583,7 @@ func TestPayload_CGO_ReaderandWriter(t *testing.T) {
assert.NotNil(t, err)
err = w.FinishPayloadWriter()
assert.NotNil(t, err)
assert.Nil(t, err)
err = w.AddBinaryVectorToPayload([]byte{}, 8)
assert.NotNil(t, err)
@ -604,7 +604,7 @@ func TestPayload_CGO_ReaderandWriter(t *testing.T) {
defer w.Close()
err = w.FinishPayloadWriter()
assert.NotNil(t, err)
assert.Nil(t, err)
err = w.AddFloatVectorToPayload([]float32{}, 8)
assert.NotNil(t, err)

4
internal/storage/payload_test.go
View File

@ -583,7 +583,7 @@ func TestPayload_ReaderAndWriter(t *testing.T) {
assert.NotNil(t, err)
err = w.FinishPayloadWriter()
assert.NotNil(t, err)
assert.Nil(t, err)
err = w.AddBinaryVectorToPayload([]byte{}, 8)
assert.NotNil(t, err)
@ -604,7 +604,7 @@ func TestPayload_ReaderAndWriter(t *testing.T) {
defer w.Close()
err = w.FinishPayloadWriter()
assert.NotNil(t, err)
assert.Nil(t, err)
err = w.AddFloatVectorToPayload([]float32{}, 8)
assert.NotNil(t, err)