diff --git a/pkg/vm/engine/tae/containers/types.go b/pkg/vm/engine/tae/containers/types.go index 44d1acdaec32..532d22f0e9b1 100644 --- a/pkg/vm/engine/tae/containers/types.go +++ b/pkg/vm/engine/tae/containers/types.go @@ -37,6 +37,9 @@ type ItOpT[T any] func(v T, isNull bool, row int) error type Vector interface { GetType() *types.Type + IsConst() bool + IsConstNull() bool + // Deep copy ops Get(i int) any Append(v any, isNull bool) diff --git a/pkg/vm/engine/tae/containers/vector.go b/pkg/vm/engine/tae/containers/vector.go index 1d9f3fbb0806..53208d67d97a 100644 --- a/pkg/vm/engine/tae/containers/vector.go +++ b/pkg/vm/engine/tae/containers/vector.go @@ -42,6 +42,61 @@ type vectorWrapper struct { element *vectorPoolElement } +func NewConstFixed[T any](typ types.Type, val T, length int, opts ...Options) *vectorWrapper { + var ( + alloc *mpool.MPool + ) + if len(opts) > 0 { + alloc = opts[0].Allocator + } + if alloc == nil { + alloc = common.DefaultAllocator + } + vec := &vectorWrapper{ + wrapped: vector.NewConstFixed[T](typ, val, length, alloc), + } + vec.mpool = alloc + return vec +} + +func NewConstBytes(typ types.Type, val []byte, length int, opts ...Options) *vectorWrapper { + var ( + alloc *mpool.MPool + ) + if len(opts) > 0 { + alloc = opts[0].Allocator + } + if alloc == nil { + alloc = common.DefaultAllocator + } + vec := &vectorWrapper{ + wrapped: vector.NewConstBytes(typ, val, length, alloc), + } + vec.mpool = alloc + return vec +} + +func NewConstNullVector( + typ types.Type, + length int, + opts ...Options, +) *vectorWrapper { + vec := &vectorWrapper{ + wrapped: vector.NewConstNull(typ, length, nil), + } + var ( + alloc *mpool.MPool + ) + if len(opts) > 0 { + alloc = opts[0].Allocator + } + if alloc == nil { + alloc = common.DefaultAllocator + } + vec.mpool = alloc + return vec +} + func NewVector(typ types.Type, opts ...Options) *vectorWrapper { vec := &vectorWrapper{ wrapped: vector.NewVec(typ), @@ -217,6 +272,14 @@ func (vec *vectorWrapper) NullCount() int { return 0 } +func (vec *vectorWrapper) IsConst() bool { + return vec.wrapped.IsConst() +} + +func (vec *vectorWrapper) IsConstNull() bool { + return vec.wrapped.IsConstNull() +} + // conver a const vectorWrapper to a normal one, getting ready to edit func (vec *vectorWrapper) TryConvertConst() Vector { if vec.wrapped.IsConstNull() { diff --git a/pkg/vm/engine/tae/containers/vector_test.go b/pkg/vm/engine/tae/containers/vector_test.go index 1c88f4fe23a1..adcb380e9d02 100644 --- a/pkg/vm/engine/tae/containers/vector_test.go +++ b/pkg/vm/engine/tae/containers/vector_test.go @@ -751,3 +751,82 @@ func TestVectorPool3(t *testing.T) { vec1.Close() vec2.Close() } + +func TestConstNullVector(t *testing.T) { + vec := NewConstNullVector(types.T_int32.ToType(), 10) + defer vec.Close() + assert.Equal(t, 10, vec.Length()) + assert.True(t, vec.IsConstNull()) + + for i := 0; i < vec.Length(); i++ { + assert.True(t, vec.IsNull(i)) + } + + var w bytes.Buffer + _, err := vec.WriteTo(&w) + assert.NoError(t, err) + + vec2 := MakeVector(types.T_int32.ToType()) + defer vec2.Close() + _, err = vec2.ReadFrom(&w) + assert.NoError(t, err) + assert.Equal(t, 10, vec2.Length()) + assert.True(t, vec2.IsConstNull()) + for i := 0; i < vec2.Length(); i++ { + assert.True(t, vec2.IsNull(i)) + } +} + +func TestConstVector(t *testing.T) { + vec := NewConstFixed[int32](types.T_int32.ToType(), int32(1), 10) + defer vec.Close() + + assert.Equal(t, 10, vec.Length()) + assert.True(t, vec.IsConst()) + assert.False(t, vec.IsConstNull()) + + for i := 0; i < vec.Length(); i++ { + v := vec.Get(i).(int32) + assert.Equal(t, int32(1), v) + } + + vec2 := NewConstBytes(types.T_char.ToType(), []byte("abc"), 10) + defer vec2.Close() + assert.Equal(t, 10, vec2.Length()) + assert.True(t, vec2.IsConst()) + assert.False(t, vec2.IsConstNull()) + + for i := 0; i < vec2.Length(); i++ { + assert.Equal(t, []byte("abc"), vec2.Get(i).([]byte)) + } + + var w bytes.Buffer + _, err := vec.WriteTo(&w) + assert.NoError(t, err) + + vec3 := MakeVector(types.T_int32.ToType()) + defer vec3.Close() + _, err = vec3.ReadFrom(&w) + assert.NoError(t, err) + assert.True(t, vec3.IsConst()) + assert.False(t, vec3.IsConstNull()) + assert.Equal(t, 10, vec3.Length()) + for i := 0; i < vec3.Length(); i++ { + assert.Equal(t, int32(1), vec3.Get(i).(int32)) + } + + w.Reset() + _, err = vec2.WriteTo(&w) + assert.NoError(t, err) + + vec4 := MakeVector(types.T_char.ToType()) + defer vec4.Close() + _, err = vec4.ReadFrom(&w) + assert.NoError(t, err) + assert.True(t, vec4.IsConst()) + assert.False(t, vec4.IsConstNull()) + assert.Equal(t, 10, vec4.Length()) + for i := 0; i < vec4.Length(); i++ { + assert.Equal(t, []byte("abc"), vec4.Get(i).([]byte)) + } +}