Skip to content

Commit 95ed5df

Browse files
authored
fix: Put null helpers back (#1002)
Return of the code from #913
1 parent 466796b commit 95ed5df

File tree

2 files changed

+85
-6
lines changed

2 files changed

+85
-6
lines changed

plugin/nulls.go

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
package plugin
2+
3+
import (
4+
"github.com/apache/arrow/go/v13/arrow"
5+
"github.com/apache/arrow/go/v13/arrow/array"
6+
"github.com/apache/arrow/go/v13/arrow/memory"
7+
)
8+
9+
// TODO(v4): use in v4
10+
//
11+
// nolint:unused
12+
func stripNullsFromLists(records []arrow.Record) {
13+
for i := range records {
14+
cols := records[i].Columns()
15+
for c, col := range cols {
16+
if col.DataType().ID() != arrow.LIST {
17+
continue
18+
}
19+
20+
list := col.(*array.List)
21+
bldr := array.NewListBuilder(memory.DefaultAllocator, list.DataType().(*arrow.ListType).Elem())
22+
for j := 0; j < list.Len(); j++ {
23+
if list.IsNull(j) {
24+
bldr.AppendNull()
25+
continue
26+
}
27+
bldr.Append(true)
28+
vBldr := bldr.ValueBuilder()
29+
from, to := list.ValueOffsets(j)
30+
slc := array.NewSlice(list.ListValues(), from, to)
31+
for k := 0; k < int(to-from); k++ {
32+
if slc.IsNull(k) {
33+
continue
34+
}
35+
err := vBldr.AppendValueFromString(slc.ValueStr(k))
36+
if err != nil {
37+
panic(err)
38+
}
39+
}
40+
}
41+
cols[c] = bldr.NewArray()
42+
}
43+
records[i] = array.NewRecord(records[i].Schema(), cols, records[i].NumRows())
44+
}
45+
}
46+
47+
type AllowNullFunc func(arrow.DataType) bool
48+
49+
// TODO(v4): use in v4
50+
//
51+
// nolint:unused
52+
func (f AllowNullFunc) replaceNullsByEmpty(records []arrow.Record) {
53+
if f == nil {
54+
return
55+
}
56+
for i := range records {
57+
cols := records[i].Columns()
58+
for c, col := range records[i].Columns() {
59+
if col.NullN() == 0 || f(col.DataType()) {
60+
continue
61+
}
62+
63+
builder := array.NewBuilder(memory.DefaultAllocator, records[i].Column(c).DataType())
64+
for j := 0; j < col.Len(); j++ {
65+
if col.IsNull(j) {
66+
builder.AppendEmptyValue()
67+
continue
68+
}
69+
70+
if err := builder.AppendValueFromString(col.ValueStr(j)); err != nil {
71+
panic(err)
72+
}
73+
}
74+
cols[c] = builder.NewArray()
75+
}
76+
records[i] = array.NewRecord(records[i].Schema(), cols, records[i].NumRows())
77+
}
78+
}

plugin/testing_write.go

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"context"
55
"testing"
66

7+
"github.com/apache/arrow/go/v13/arrow"
78
"github.com/cloudquery/plugin-sdk/v4/schema"
89
)
910

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

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

5758
type NewPluginFunc func() *Plugin
5859

59-
// func WithTestSourceAllowNull(allowNull func(arrow.DataType) bool) func(o *WriterTestSuite) {
60-
// return func(o *WriterTestSuite) {
61-
// o.allowNull = allowNull
62-
// }
63-
// }
60+
func WithTestSourceAllowNull(allowNull func(arrow.DataType) bool) func(o *WriterTestSuite) {
61+
return func(o *WriterTestSuite) {
62+
o.allowNull = allowNull
63+
}
64+
}
6465

6566
func WithTestIgnoreNullsInLists() func(o *WriterTestSuite) {
6667
return func(o *WriterTestSuite) {

0 commit comments

Comments
 (0)