dn vector wrapper supports const vector (#11133)

DN vector wrapper supports const vector

Approved by: @LeftHandCold

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)

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() {

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))
+	}
+}