From c02763d49a1d3dfb38cf395e0fca40eecd36dc03 Mon Sep 17 00:00:00 2001
From: congqixia <congqi.xia@zilliz.com>
Date: Tue, 4 Mar 2025 20:58:04 +0800
Subject: [PATCH] fix: [GoSDK] Pass only valid data for nullable column
 (#40328)

Related to #40327

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
---
 client/column/array.go         |  1 +
 client/column/generic_base.go  |  7 +++++--
 client/column/nullable.go      | 13 ++++++++++---
 client/column/nullable_test.go | 14 +++++++-------
 client/entity/field.go         |  1 +
 5 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/client/column/array.go b/client/column/array.go
index 90251fe689..4f25e41793 100644
--- a/client/column/array.go
+++ b/client/column/array.go
@@ -32,6 +32,7 @@ func (c *columnArrayBase[T]) FieldData() *schemapb.FieldData {
 	fd := &schemapb.FieldData{
 		Type:      schemapb.DataType_Array,
 		FieldName: c.name,
+		ValidData: c.validData,
 	}
 
 	data := make([]*schemapb.ScalarField, 0, c.Len())
diff --git a/client/column/generic_base.go b/client/column/generic_base.go
index e2e001828f..1d5021a865 100644
--- a/client/column/generic_base.go
+++ b/client/column/generic_base.go
@@ -49,6 +49,9 @@ func (c *genericColumnBase[T]) Type() entity.FieldType {
 }
 
 func (c *genericColumnBase[T]) Len() int {
+	if c.validData != nil {
+		return len(c.validData)
+	}
 	return len(c.values)
 }
 
@@ -166,9 +169,9 @@ func (c *genericColumnBase[T]) AppendNull() error {
 	if !c.nullable {
 		return errors.New("append null to not nullable column")
 	}
-	var v T
+	// var v T
 	c.validData = append(c.validData, true)
-	c.values = append(c.values, v)
+	// c.values = append(c.values, v)
 	return nil
 }
 
diff --git a/client/column/nullable.go b/client/column/nullable.go
index 3ff68d083b..52f2b80e32 100644
--- a/client/column/nullable.go
+++ b/client/column/nullable.go
@@ -16,7 +16,10 @@
 
 package column
 
-import "github.com/cockroachdb/errors"
+import (
+	"github.com/cockroachdb/errors"
+	"github.com/samber/lo"
+)
 
 var (
 	// scalars
@@ -55,9 +58,13 @@ type NullableColumnCreator[col interface {
 
 func (c NullableColumnCreator[col, T]) New(name string, values []T, validData []bool) (col, error) {
 	var result col
-	if len(values) != len(validData) {
-		return result, errors.New("values & validData slice has different length")
+	validCnt := lo.CountBy(validData, func(v bool) bool {
+		return v
+	})
+	if validCnt != len(values) {
+		return result, errors.Newf("values number(%d) does not match valid count(%d)", len(values), validCnt)
 	}
+
 	result = c.base(name, values)
 	result.withValidData(validData)
 
diff --git a/client/column/nullable_test.go b/client/column/nullable_test.go
index eec2cbb489..c22ede658b 100644
--- a/client/column/nullable_test.go
+++ b/client/column/nullable_test.go
@@ -33,7 +33,7 @@ type NullableScalarSuite struct {
 func (s *NullableScalarSuite) TestBasic() {
 	s.Run("nullable_bool", func() {
 		name := fmt.Sprintf("field_%d", rand.Intn(1000))
-		data := []bool{true, false}
+		data := []bool{false}
 		validData := []bool{true, false}
 		column, err := NewNullableColumnBool(name, data, validData)
 		s.NoError(err)
@@ -63,7 +63,7 @@ func (s *NullableScalarSuite) TestBasic() {
 
 	s.Run("nullable_int8", func() {
 		name := fmt.Sprintf("field_%d", rand.Intn(1000))
-		data := []int8{1, 2, 3}
+		data := []int8{1, 3}
 		validData := []bool{true, false, true}
 		column, err := NewNullableColumnInt8(name, data, validData)
 		s.NoError(err)
@@ -93,7 +93,7 @@ func (s *NullableScalarSuite) TestBasic() {
 
 	s.Run("nullable_int16", func() {
 		name := fmt.Sprintf("field_%d", rand.Intn(1000))
-		data := []int16{1, 2, 3}
+		data := []int16{1, 3}
 		validData := []bool{true, false, true}
 		column, err := NewNullableColumnInt16(name, data, validData)
 		s.NoError(err)
@@ -123,7 +123,7 @@ func (s *NullableScalarSuite) TestBasic() {
 
 	s.Run("nullable_int32", func() {
 		name := fmt.Sprintf("field_%d", rand.Intn(1000))
-		data := []int32{1, 2, 3}
+		data := []int32{1, 3}
 		validData := []bool{true, false, true}
 		column, err := NewNullableColumnInt32(name, data, validData)
 		s.NoError(err)
@@ -153,7 +153,7 @@ func (s *NullableScalarSuite) TestBasic() {
 
 	s.Run("nullable_int64", func() {
 		name := fmt.Sprintf("field_%d", rand.Intn(1000))
-		data := []int64{1, 2, 3}
+		data := []int64{1, 3}
 		validData := []bool{true, false, true}
 		column, err := NewNullableColumnInt64(name, data, validData)
 		s.NoError(err)
@@ -183,7 +183,7 @@ func (s *NullableScalarSuite) TestBasic() {
 
 	s.Run("nullable_float", func() {
 		name := fmt.Sprintf("field_%d", rand.Intn(1000))
-		data := []float32{0.1, 0.2, 0.3}
+		data := []float32{0.1, 0.3}
 		validData := []bool{true, false, true}
 		column, err := NewNullableColumnFloat(name, data, validData)
 		s.NoError(err)
@@ -213,7 +213,7 @@ func (s *NullableScalarSuite) TestBasic() {
 
 	s.Run("nullable_double", func() {
 		name := fmt.Sprintf("field_%d", rand.Intn(1000))
-		data := []float64{0.1, 0.2, 0.3}
+		data := []float64{0.1, 0.3}
 		validData := []bool{true, false, true}
 		column, err := NewNullableColumnDouble(name, data, validData)
 		s.NoError(err)
diff --git a/client/entity/field.go b/client/entity/field.go
index 3ff3f7e4f2..37cd8c93e9 100644
--- a/client/entity/field.go
+++ b/client/entity/field.go
@@ -220,6 +220,7 @@ func (f *Field) ProtoMessage() *schemapb.FieldSchema {
 		IsPartitionKey:  f.IsPartitionKey,
 		IsClusteringKey: f.IsClusteringKey,
 		ElementType:     schemapb.DataType(f.ElementType),
+		Nullable:        f.Nullable,
 	}
 }