Skip to content

Commit

Permalink
put null helpers back
Browse files Browse the repository at this point in the history
  • Loading branch information
candiduslynx committed Jun 29, 2023
1 parent d1b2b73 commit c030939
Show file tree
Hide file tree
Showing 2 changed files with 79 additions and 6 deletions.
72 changes: 72 additions & 0 deletions plugin/nulls.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
package plugin

import (
"github.com/apache/arrow/go/v13/arrow"
"github.com/apache/arrow/go/v13/arrow/array"
"github.com/apache/arrow/go/v13/arrow/memory"
)

func stripNullsFromLists(records []arrow.Record) {
for i := range records {
cols := records[i].Columns()
for c, col := range cols {
if col.DataType().ID() != arrow.LIST {
continue
}

list := col.(*array.List)
bldr := array.NewListBuilder(memory.DefaultAllocator, list.DataType().(*arrow.ListType).Elem())
for j := 0; j < list.Len(); j++ {
if list.IsNull(j) {
bldr.AppendNull()
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)
}
}
}
cols[c] = bldr.NewArray()
}
records[i] = array.NewRecord(records[i].Schema(), cols, records[i].NumRows())
}
}

type AllowNullFunc func(arrow.DataType) bool

func (f AllowNullFunc) replaceNullsByEmpty(records []arrow.Record) {
if f == nil {
return
}
for i := range records {
cols := records[i].Columns()
for c, col := range records[i].Columns() {
if col.NullN() == 0 || f(col.DataType()) {
continue
}

builder := array.NewBuilder(memory.DefaultAllocator, records[i].Column(c).DataType())
for j := 0; j < col.Len(); j++ {
if col.IsNull(j) {
builder.AppendEmptyValue()
continue
}

if err := builder.AppendValueFromString(col.ValueStr(j)); err != nil {
panic(err)
}
}
cols[c] = builder.NewArray()
}
records[i] = array.NewRecord(records[i].Schema(), cols, records[i].NumRows())
}
}
13 changes: 7 additions & 6 deletions plugin/testing_write.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"context"
"testing"

"github.com/apache/arrow/go/v13/arrow"
"github.com/cloudquery/plugin-sdk/v4/schema"
)

Expand All @@ -16,7 +17,7 @@ type WriterTestSuite struct {
// Destinations that have problems representing some data types should provide a custom implementation here.
// If this param is empty, the default is to allow all data types to be nullable.
// When the value returned by this func is `true` the comparison is made with the empty value instead of null.
// allowNull AllowNullFunc
allowNull AllowNullFunc

// IgnoreNullsInLists allows stripping null values from lists before comparison.
// Destination setups that don't support nulls in lists should set this to true.
Expand Down Expand Up @@ -56,11 +57,11 @@ type WriterTestSuiteTests struct {

type NewPluginFunc func() *Plugin

// func WithTestSourceAllowNull(allowNull func(arrow.DataType) bool) func(o *WriterTestSuite) {
// return func(o *WriterTestSuite) {
// o.allowNull = allowNull
// }
// }
func WithTestSourceAllowNull(allowNull func(arrow.DataType) bool) func(o *WriterTestSuite) {
return func(o *WriterTestSuite) {
o.allowNull = allowNull
}
}

func WithTestIgnoreNullsInLists() func(o *WriterTestSuite) {
return func(o *WriterTestSuite) {
Expand Down

0 comments on commit c030939

Please sign in to comment.