fix: Process nulls for tested types, too (maps, lists, structs) (#1125)

Author: candiduslynx

plugin/nulls.go

@@ -6,75 +6,108 @@ import (
 	"github.com/apache/arrow/go/v13/arrow/memory"
 )
 
-func stripNullsFromLists(record arrow.Record) arrow.Record {
-	cols := record.Columns()
-	for c, col := range cols {
-		list, ok := col.(array.ListLike)
-		if !ok {
-			continue
-		}
-		if _, ok := list.(*array.Map); ok {
-			// maps also correspond to array.ListLike
+func stripNullsFromLists(list array.ListLike) array.ListLike {
+	// TODO: handle Arrow maps separately if required
+
+	if list.NullN() == 0 {
+		return list
+	}
+
+	bldr := array.NewBuilder(memory.DefaultAllocator, list.DataType()).(array.ListLikeBuilder)
+	for j := 0; j < list.Len(); j++ {
+		if list.IsNull(j) {
+			bldr.AppendNull()
 			continue
 		}
-
-		bldr := array.NewListBuilder(memory.DefaultAllocator, list.DataType().(arrow.ListLikeType).Elem())
-		for j := 0; j < list.Len(); j++ {
-			if list.IsNull(j) {
-				bldr.AppendNull()
+		bldr.Append(true)
+		vBldr := bldr.ValueBuilder()
+		from, to := list.ValueOffsets(j)
+		slc := array.NewSlice(list.ListValues(), from, to)
+		for k := 0; k < int(to-from); k++ {
+			if slc.IsNull(k) {
 				continue
 			}
-			bldr.Append(true)
-			vBldr := bldr.ValueBuilder()
-			from, to := list.ValueOffsets(j)
-			slc := array.NewSlice(list.ListValues(), from, to)
-			for k := 0; k < int(to-from); k++ {
-				if slc.IsNull(k) {
-					continue
-				}
-				err := vBldr.AppendValueFromString(slc.ValueStr(k))
-				if err != nil {
-					panic(err)
-				}
+			err := vBldr.AppendValueFromString(slc.ValueStr(k))
+			if err != nil {
+				panic(err)
 			}
 		}
-		cols[c] = bldr.NewArray()
 	}
-	return array.NewRecord(record.Schema(), cols, record.NumRows())
+
+	return bldr.NewArray().(array.ListLike)
 }
 
 type AllowNullFunc func(arrow.DataType) bool
 
-func (s *WriterTestSuite) replaceNullsByEmpty(record arrow.Record) arrow.Record {
+func (s *WriterTestSuite) replaceNullsByEmpty(arr arrow.Array) arrow.Array {
 	if s.allowNull == nil {
-		return record
+		return arr
 	}
 
-	cols := record.Columns()
-	for c, col := range cols {
-		if col.NullN() == 0 || s.allowNull(col.DataType()) {
-			continue
-		}
-
-		builder := array.NewBuilder(memory.DefaultAllocator, col.DataType())
-		for j := 0; j < col.Len(); j++ {
-			if col.IsNull(j) {
+	if !s.allowNull(arr.DataType()) && arr.NullN() > 0 {
+		builder := array.NewBuilder(memory.DefaultAllocator, arr.DataType())
+		for j := 0; j < arr.Len(); j++ {
+			if arr.IsNull(j) {
 				builder.AppendEmptyValue()
 				continue
 			}
 
-			if err := builder.AppendValueFromString(col.ValueStr(j)); err != nil {
+			if err := builder.AppendValueFromString(arr.ValueStr(j)); err != nil {
 				panic(err)
 			}
 		}
-		cols[c] = builder.NewArray()
+
+		arr = builder.NewArray()
+	}
+
+	// we need to process the nested arrays, too
+	return s.replaceNullsByEmptyNestedArray(arr)
+}
+
+func (s *WriterTestSuite) replaceNullsByEmptyNestedArray(arr arrow.Array) arrow.Array {
+	if s.allowNull == nil {
+		return arr
+	}
+
+	switch arr := arr.(type) {
+	case array.ListLike: // TODO: handle Arrow maps separately if required
+		values := s.handleNullsArray(arr.ListValues())
+		return array.MakeFromData(
+			array.NewData(arr.DataType(), arr.Len(),
+				arr.Data().Buffers(),
+				[]arrow.ArrayData{values.Data()},
+				arr.NullN(), arr.Data().Offset(),
+			),
+		)
+	case *array.Struct:
+		children := make([]arrow.ArrayData, arr.NumField())
+		for i := 0; i < arr.NumField(); i++ {
+			children[i] = s.handleNullsArray(arr.Field(i)).Data()
+		}
+		return array.MakeFromData(
+			array.NewData(arr.DataType(), arr.Len(),
+				arr.Data().Buffers(),
+				children,
+				arr.NullN(), arr.Data().Offset(),
+			),
+		)
+	default:
+		return arr
 	}
-	return array.NewRecord(record.Schema(), cols, record.NumRows())
 }
 
 func (s *WriterTestSuite) handleNulls(record arrow.Record) arrow.Record {
-	if s.ignoreNullsInLists {
-		record = stripNullsFromLists(record)
+	cols := record.Columns()
+	for c, col := range cols {
+		cols[c] = s.handleNullsArray(col)
+	}
+	return array.NewRecord(record.Schema(), cols, record.NumRows())
+}
+
+func (s *WriterTestSuite) handleNullsArray(arr arrow.Array) arrow.Array {
+	if list, ok := arr.(array.ListLike); ok && s.ignoreNullsInLists {
+		arr = stripNullsFromLists(list) // TODO: handle Arrow maps separately if required
 	}
-	return s.replaceNullsByEmpty(record)
+
+	return s.replaceNullsByEmpty(arr)
 }

plugin/nulls_test.go

@@ -0,0 +1,62 @@
+package plugin
+
+import (
+	"github.com/apache/arrow/go/v13/arrow"
+	"github.com/apache/arrow/go/v13/arrow/array"
+	"github.com/cloudquery/plugin-sdk/v4/schema"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"testing"
+	"time"
+)
+
+func TestWithTestSourceAllowNull(t *testing.T) {
+	s := &WriterTestSuite{allowNull: func(dt arrow.DataType) bool {
+		switch dt.(type) {
+		case *arrow.StructType, arrow.ListLikeType:
+			return false
+		default:
+			return true
+		}
+	}}
+
+	tg := schema.NewTestDataGenerator()
+	source := schema.TestTable("allow_null", schema.TestSourceOptions{})
+	resource := s.handleNulls(tg.Generate(source, schema.GenTestDataOptions{
+		SourceName: "allow_null",
+		SyncTime:   time.Now(),
+		MaxRows:    100,
+		NullRows:   false,
+	})[0])
+	for _, c := range resource.Columns() {
+		assertNoNulls(t, s.allowNull, c)
+	}
+
+	resource = s.handleNulls(tg.Generate(source, schema.GenTestDataOptions{
+		SourceName: "allow_null",
+		SyncTime:   time.Now(),
+		MaxRows:    100,
+		NullRows:   true,
+	})[0])
+	for _, c := range resource.Columns() {
+		assertNoNulls(t, s.allowNull, c)
+	}
+}
+
+func assertNoNulls(t *testing.T, allowNull AllowNullFunc, arr arrow.Array) {
+	require.NotNil(t, allowNull)
+
+	if !allowNull(arr.DataType()) {
+		assert.Zero(t, arr.NullN())
+	}
+
+	// traverse
+	switch arr := arr.(type) {
+	case array.ListLike:
+		assertNoNulls(t, allowNull, arr.ListValues())
+	case *array.Struct:
+		for i := 0; i < arr.NumField(); i++ {
+			assertNoNulls(t, allowNull, arr.Field(i))
+		}
+	}
+}