mirror of https://github.com/milvus-io/milvus.git
neza2017
yefu.chen
parent
881be7f3e3
commit
5512ce8d8a
|
|
@ -0,0 +1,283 @@
|
|||
package storage
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/binary"
|
||||
"testing"
|
||||
"time"
|
||||
"unsafe"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/zilliztech/milvus-distributed/internal/proto/schemapb"
|
||||
"github.com/zilliztech/milvus-distributed/internal/util/tsoutil"
|
||||
)
|
||||
|
||||
func checkEventHeader(
|
||||
t *testing.T,
|
||||
buf []byte,
|
||||
tc EventTypeCode,
|
||||
svrID int32,
|
||||
length int32) {
|
||||
ts := UnsafeReadInt64(buf, 0)
|
||||
assert.Greater(t, ts, int64(0))
|
||||
curts := time.Now().UnixNano() / int64(time.Millisecond)
|
||||
curts = int64(tsoutil.ComposeTS(curts, 0))
|
||||
assert.GreaterOrEqual(t, curts, ts)
|
||||
utc := UnsafeReadInt8(buf, int(unsafe.Sizeof(ts)))
|
||||
assert.Equal(t, EventTypeCode(utc), tc)
|
||||
usid := UnsafeReadInt32(buf, int(unsafe.Sizeof(ts)+unsafe.Sizeof(utc)))
|
||||
assert.Equal(t, usid, svrID)
|
||||
elen := UnsafeReadInt32(buf, int(unsafe.Sizeof(ts)+unsafe.Sizeof(utc)+unsafe.Sizeof(usid)))
|
||||
assert.Equal(t, elen, length)
|
||||
npos := UnsafeReadInt32(buf, int(unsafe.Sizeof(ts)+unsafe.Sizeof(utc)+unsafe.Sizeof(usid)+unsafe.Sizeof(elen)))
|
||||
assert.Equal(t, npos, length)
|
||||
}
|
||||
|
||||
func TestEventWriterAndReader(t *testing.T) {
|
||||
insertT := func(t *testing.T,
|
||||
dt schemapb.DataType,
|
||||
ir1 func(w *insertEventWriter) error,
|
||||
ir2 func(w *insertEventWriter) error,
|
||||
iw func(w *insertEventWriter) error,
|
||||
ev interface{},
|
||||
) {
|
||||
w, err := newInsertEventWriter(dt, 0)
|
||||
assert.Nil(t, err)
|
||||
w.SetStartTimestamp(tsoutil.ComposeTS(10, 0))
|
||||
w.SetEndTimestamp(tsoutil.ComposeTS(100, 0))
|
||||
err = ir1(w)
|
||||
assert.Nil(t, err)
|
||||
err = iw(w)
|
||||
assert.NotNil(t, err)
|
||||
err = ir2(w)
|
||||
assert.Nil(t, err)
|
||||
err = w.Finish()
|
||||
assert.Nil(t, err)
|
||||
|
||||
var buf bytes.Buffer
|
||||
err = w.Write(&buf)
|
||||
assert.Nil(t, err)
|
||||
err = w.Close()
|
||||
assert.Nil(t, err)
|
||||
|
||||
wBuf := buf.Bytes()
|
||||
checkEventHeader(t, wBuf, InsertEventType, ServerID, int32(len(wBuf)))
|
||||
st := UnsafeReadInt64(wBuf, binary.Size(eventHeader{}))
|
||||
assert.Equal(t, Timestamp(st), tsoutil.ComposeTS(10, 0))
|
||||
et := UnsafeReadInt64(wBuf, binary.Size(eventHeader{})+int(unsafe.Sizeof(st)))
|
||||
assert.Equal(t, Timestamp(et), tsoutil.ComposeTS(100, 0))
|
||||
|
||||
payloadOffset := binary.Size(eventHeader{}) + binary.Size(insertEventData{})
|
||||
pBuf := wBuf[payloadOffset:]
|
||||
pR, err := NewPayloadReader(dt, pBuf)
|
||||
assert.Nil(t, err)
|
||||
vals, _, err := pR.GetDataFromPayload()
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, vals, ev)
|
||||
err = pR.Close()
|
||||
assert.Nil(t, err)
|
||||
|
||||
r, err := newEventReader(dt, bytes.NewBuffer(wBuf))
|
||||
assert.Nil(t, err)
|
||||
payload, _, err := r.GetDataFromPayload()
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, payload, ev)
|
||||
|
||||
err = r.Close()
|
||||
assert.Nil(t, err)
|
||||
}
|
||||
|
||||
t.Run("insert_bool", func(t *testing.T) {
|
||||
insertT(t, schemapb.DataType_BOOL,
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]bool{true, false, true})
|
||||
},
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]bool{false, true, false})
|
||||
},
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]int{1, 2, 3, 4, 5})
|
||||
},
|
||||
[]bool{true, false, true, false, true, false})
|
||||
})
|
||||
|
||||
t.Run("insert_int8", func(t *testing.T) {
|
||||
insertT(t, schemapb.DataType_INT8,
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]int8{1, 2, 3})
|
||||
},
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]int8{4, 5, 6})
|
||||
},
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]int{1, 2, 3, 4, 5})
|
||||
},
|
||||
[]int8{1, 2, 3, 4, 5, 6})
|
||||
})
|
||||
|
||||
t.Run("insert_int16", func(t *testing.T) {
|
||||
insertT(t, schemapb.DataType_INT16,
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]int16{1, 2, 3})
|
||||
},
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]int16{4, 5, 6})
|
||||
},
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]int{1, 2, 3, 4, 5})
|
||||
},
|
||||
[]int16{1, 2, 3, 4, 5, 6})
|
||||
})
|
||||
|
||||
t.Run("insert_int32", func(t *testing.T) {
|
||||
insertT(t, schemapb.DataType_INT32,
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]int32{1, 2, 3})
|
||||
},
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]int32{4, 5, 6})
|
||||
},
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]int{1, 2, 3, 4, 5})
|
||||
},
|
||||
[]int32{1, 2, 3, 4, 5, 6})
|
||||
})
|
||||
|
||||
t.Run("insert_int64", func(t *testing.T) {
|
||||
insertT(t, schemapb.DataType_INT64,
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]int64{1, 2, 3})
|
||||
},
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]int64{4, 5, 6})
|
||||
},
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]int{1, 2, 3, 4, 5})
|
||||
},
|
||||
[]int64{1, 2, 3, 4, 5, 6})
|
||||
})
|
||||
|
||||
t.Run("insert_float32", func(t *testing.T) {
|
||||
insertT(t, schemapb.DataType_FLOAT,
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]float32{1, 2, 3})
|
||||
},
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]float32{4, 5, 6})
|
||||
},
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]int{1, 2, 3, 4, 5})
|
||||
},
|
||||
[]float32{1, 2, 3, 4, 5, 6})
|
||||
})
|
||||
|
||||
t.Run("insert_float64", func(t *testing.T) {
|
||||
insertT(t, schemapb.DataType_DOUBLE,
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]float64{1, 2, 3})
|
||||
},
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]float64{4, 5, 6})
|
||||
},
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]int{1, 2, 3, 4, 5})
|
||||
},
|
||||
[]float64{1, 2, 3, 4, 5, 6})
|
||||
})
|
||||
|
||||
t.Run("insert_binary_vector", func(t *testing.T) {
|
||||
insertT(t, schemapb.DataType_VECTOR_BINARY,
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]byte{1, 2, 3, 4}, 16)
|
||||
},
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]byte{5, 6, 7, 8}, 16)
|
||||
},
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]int{1, 2, 3, 4, 5, 6}, 16)
|
||||
},
|
||||
[]byte{1, 2, 3, 4, 5, 6, 7, 8})
|
||||
})
|
||||
|
||||
t.Run("insert_float_vector", func(t *testing.T) {
|
||||
insertT(t, schemapb.DataType_VECTOR_FLOAT,
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]float32{1, 2, 3, 4}, 2)
|
||||
},
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]float32{5, 6, 7, 8}, 2)
|
||||
},
|
||||
func(w *insertEventWriter) error {
|
||||
return w.AddDataToPayload([]int{1, 2, 3, 4, 5, 6}, 2)
|
||||
},
|
||||
[]float32{1, 2, 3, 4, 5, 6, 7, 8})
|
||||
})
|
||||
|
||||
t.Run("insert_string", func(t *testing.T) {
|
||||
w, err := newInsertEventWriter(schemapb.DataType_STRING, 0)
|
||||
assert.Nil(t, err)
|
||||
w.SetStartTimestamp(tsoutil.ComposeTS(10, 0))
|
||||
w.SetEndTimestamp(tsoutil.ComposeTS(100, 0))
|
||||
err = w.AddDataToPayload("1234")
|
||||
assert.Nil(t, err)
|
||||
err = w.AddOneStringToPayload("567890")
|
||||
assert.Nil(t, err)
|
||||
err = w.AddOneStringToPayload("abcdefg")
|
||||
assert.Nil(t, err)
|
||||
err = w.AddDataToPayload([]int{1, 2, 3})
|
||||
assert.NotNil(t, err)
|
||||
err = w.Finish()
|
||||
assert.Nil(t, err)
|
||||
|
||||
var buf bytes.Buffer
|
||||
err = w.Write(&buf)
|
||||
assert.Nil(t, err)
|
||||
err = w.Close()
|
||||
assert.Nil(t, err)
|
||||
|
||||
wBuf := buf.Bytes()
|
||||
checkEventHeader(t, wBuf, InsertEventType, ServerID, int32(len(wBuf)))
|
||||
st := UnsafeReadInt64(wBuf, binary.Size(eventHeader{}))
|
||||
assert.Equal(t, Timestamp(st), tsoutil.ComposeTS(10, 0))
|
||||
et := UnsafeReadInt64(wBuf, binary.Size(eventHeader{})+int(unsafe.Sizeof(st)))
|
||||
assert.Equal(t, Timestamp(et), tsoutil.ComposeTS(100, 0))
|
||||
|
||||
payloadOffset := binary.Size(eventHeader{}) + binary.Size(insertEventData{})
|
||||
pBuf := wBuf[payloadOffset:]
|
||||
pR, err := NewPayloadReader(schemapb.DataType_STRING, pBuf)
|
||||
assert.Nil(t, err)
|
||||
|
||||
s0, err := pR.GetOneStringFromPayload(0)
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, s0, "1234")
|
||||
|
||||
s1, err := pR.GetOneStringFromPayload(1)
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, s1, "567890")
|
||||
|
||||
s2, err := pR.GetOneStringFromPayload(2)
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, s2, "abcdefg")
|
||||
|
||||
err = pR.Close()
|
||||
assert.Nil(t, err)
|
||||
|
||||
r, err := newEventReader(schemapb.DataType_STRING, bytes.NewBuffer(wBuf))
|
||||
assert.Nil(t, err)
|
||||
|
||||
s0, err = r.GetOneStringFromPayload(0)
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, s0, "1234")
|
||||
|
||||
s1, err = r.GetOneStringFromPayload(1)
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, s1, "567890")
|
||||
|
||||
s2, err = r.GetOneStringFromPayload(2)
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, s2, "abcdefg")
|
||||
|
||||
err = r.Close()
|
||||
assert.Nil(t, err)
|
||||
})
|
||||
}
|
||||
|
|
@ -5,9 +5,8 @@ import (
|
|||
"encoding/binary"
|
||||
"testing"
|
||||
|
||||
"github.com/zilliztech/milvus-distributed/internal/proto/schemapb"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/zilliztech/milvus-distributed/internal/proto/schemapb"
|
||||
)
|
||||
|
||||
func TestSizeofStruct(t *testing.T) {
|
||||
|
|
@ -34,12 +33,13 @@ func TestSizeofStruct(t *testing.T) {
|
|||
func TestEventWriter(t *testing.T) {
|
||||
insertEvent, err := newInsertEventWriter(schemapb.DataType_INT32, 0)
|
||||
assert.Nil(t, err)
|
||||
defer insertEvent.Close()
|
||||
err = insertEvent.Close()
|
||||
assert.Nil(t, err)
|
||||
|
||||
insertEvent, err = newInsertEventWriter(schemapb.DataType_INT32, 0)
|
||||
assert.Nil(t, err)
|
||||
defer insertEvent.Close()
|
||||
|
||||
err = insertEvent.AddInt64ToPayload([]int64{1, 1})
|
||||
assert.NotNil(t, err)
|
||||
err = insertEvent.AddInt32ToPayload([]int32{1, 2, 3})
|
||||
|
|
|
|||
|
|
@ -412,7 +412,10 @@ func NewPayloadReader(colType schemapb.DataType, buf []byte) (*PayloadReader, er
|
|||
if len(buf) == 0 {
|
||||
return nil, errors.New("create Payload reader failed, buffer is empty")
|
||||
}
|
||||
r := C.NewPayloadReader(C.int(colType), (*C.uchar)(unsafe.Pointer(&buf[0])), C.long(len(buf)))
|
||||
r := C.NewPayloadReader(C.int(colType), (*C.uint8_t)(unsafe.Pointer(&buf[0])), C.long(len(buf)))
|
||||
if r == nil {
|
||||
return nil, errors.New("failed to read parquet from buffer")
|
||||
}
|
||||
return &PayloadReader{payloadReaderPtr: r, colType: colType}, nil
|
||||
}
|
||||
|
||||
|
|
@ -487,6 +490,10 @@ func (r *PayloadReader) ReleasePayloadReader() error {
|
|||
}
|
||||
|
||||
func (r *PayloadReader) GetBoolFromPayload() ([]bool, error) {
|
||||
if r.colType != schemapb.DataType_BOOL {
|
||||
return nil, errors.New("incorrect data type")
|
||||
}
|
||||
|
||||
var cMsg *C.bool
|
||||
var cSize C.int
|
||||
|
||||
|
|
@ -503,6 +510,10 @@ func (r *PayloadReader) GetBoolFromPayload() ([]bool, error) {
|
|||
}
|
||||
|
||||
func (r *PayloadReader) GetInt8FromPayload() ([]int8, error) {
|
||||
if r.colType != schemapb.DataType_INT8 {
|
||||
return nil, errors.New("incorrect data type")
|
||||
}
|
||||
|
||||
var cMsg *C.int8_t
|
||||
var cSize C.int
|
||||
|
||||
|
|
@ -519,6 +530,10 @@ func (r *PayloadReader) GetInt8FromPayload() ([]int8, error) {
|
|||
}
|
||||
|
||||
func (r *PayloadReader) GetInt16FromPayload() ([]int16, error) {
|
||||
if r.colType != schemapb.DataType_INT16 {
|
||||
return nil, errors.New("incorrect data type")
|
||||
}
|
||||
|
||||
var cMsg *C.int16_t
|
||||
var cSize C.int
|
||||
|
||||
|
|
@ -535,6 +550,10 @@ func (r *PayloadReader) GetInt16FromPayload() ([]int16, error) {
|
|||
}
|
||||
|
||||
func (r *PayloadReader) GetInt32FromPayload() ([]int32, error) {
|
||||
if r.colType != schemapb.DataType_INT32 {
|
||||
return nil, errors.New("incorrect data type")
|
||||
}
|
||||
|
||||
var cMsg *C.int32_t
|
||||
var cSize C.int
|
||||
|
||||
|
|
@ -551,6 +570,10 @@ func (r *PayloadReader) GetInt32FromPayload() ([]int32, error) {
|
|||
}
|
||||
|
||||
func (r *PayloadReader) GetInt64FromPayload() ([]int64, error) {
|
||||
if r.colType != schemapb.DataType_INT64 {
|
||||
return nil, errors.New("incorrect data type")
|
||||
}
|
||||
|
||||
var cMsg *C.int64_t
|
||||
var cSize C.int
|
||||
|
||||
|
|
@ -567,6 +590,10 @@ func (r *PayloadReader) GetInt64FromPayload() ([]int64, error) {
|
|||
}
|
||||
|
||||
func (r *PayloadReader) GetFloatFromPayload() ([]float32, error) {
|
||||
if r.colType != schemapb.DataType_FLOAT {
|
||||
return nil, errors.New("incorrect data type")
|
||||
}
|
||||
|
||||
var cMsg *C.float
|
||||
var cSize C.int
|
||||
|
||||
|
|
@ -583,6 +610,10 @@ func (r *PayloadReader) GetFloatFromPayload() ([]float32, error) {
|
|||
}
|
||||
|
||||
func (r *PayloadReader) GetDoubleFromPayload() ([]float64, error) {
|
||||
if r.colType != schemapb.DataType_DOUBLE {
|
||||
return nil, errors.New("incorrect data type")
|
||||
}
|
||||
|
||||
var cMsg *C.double
|
||||
var cSize C.int
|
||||
|
||||
|
|
@ -599,6 +630,10 @@ func (r *PayloadReader) GetDoubleFromPayload() ([]float64, error) {
|
|||
}
|
||||
|
||||
func (r *PayloadReader) GetOneStringFromPayload(idx int) (string, error) {
|
||||
if r.colType != schemapb.DataType_STRING {
|
||||
return "", errors.New("incorrect data type")
|
||||
}
|
||||
|
||||
var cStr *C.char
|
||||
var cSize C.int
|
||||
|
||||
|
|
@ -615,6 +650,10 @@ func (r *PayloadReader) GetOneStringFromPayload(idx int) (string, error) {
|
|||
|
||||
// ,dimension, error
|
||||
func (r *PayloadReader) GetBinaryVectorFromPayload() ([]byte, int, error) {
|
||||
if r.colType != schemapb.DataType_VECTOR_BINARY {
|
||||
return nil, 0, errors.New("incorrect data type")
|
||||
}
|
||||
|
||||
var cMsg *C.uint8_t
|
||||
var cDim C.int
|
||||
var cLen C.int
|
||||
|
|
@ -634,6 +673,10 @@ func (r *PayloadReader) GetBinaryVectorFromPayload() ([]byte, int, error) {
|
|||
|
||||
// ,dimension, error
|
||||
func (r *PayloadReader) GetFloatVectorFromPayload() ([]float32, int, error) {
|
||||
if r.colType != schemapb.DataType_VECTOR_FLOAT {
|
||||
return nil, 0, errors.New("incorrect data type")
|
||||
}
|
||||
|
||||
var cMsg *C.float
|
||||
var cDim C.int
|
||||
var cLen C.int
|
||||
|
|
|
|||
|
|
@ -0,0 +1,33 @@
|
|||
package storage
|
||||
|
||||
import "unsafe"
|
||||
|
||||
func UnsafeReadInt8(buf []byte, idx int) int8 {
|
||||
ptr := unsafe.Pointer(&(buf[idx]))
|
||||
return *((*int8)(ptr))
|
||||
}
|
||||
|
||||
func UnsafeReadInt16(buf []byte, idx int) int16 {
|
||||
ptr := unsafe.Pointer(&(buf[idx]))
|
||||
return *((*int16)(ptr))
|
||||
}
|
||||
|
||||
func UnsafeReadInt32(buf []byte, idx int) int32 {
|
||||
ptr := unsafe.Pointer(&(buf[idx]))
|
||||
return *((*int32)(ptr))
|
||||
}
|
||||
|
||||
func UnsafeReadInt64(buf []byte, idx int) int64 {
|
||||
ptr := unsafe.Pointer(&(buf[idx]))
|
||||
return *((*int64)(ptr))
|
||||
}
|
||||
|
||||
func UnsafeReadFloat32(buf []byte, idx int) float32 {
|
||||
ptr := unsafe.Pointer(&(buf[idx]))
|
||||
return *((*float32)(ptr))
|
||||
}
|
||||
|
||||
func UnsafeReadFloat64(buf []byte, idx int) float64 {
|
||||
ptr := unsafe.Pointer(&(buf[idx]))
|
||||
return *((*float64)(ptr))
|
||||
}
|
||||
Loading…
Reference in New Issue