@@ -54,13 +54,14 @@ func TestWriteRead(t *testing.T) {
5454 readErr = cl .Read (byteReader , table , ch )
5555 close (ch )
5656 }()
57- received := make ([]arrow.Record , 0 , rows )
57+ received , total := make ([]arrow.Record , 0 , rows ), 0
5858 for got := range ch {
5959 received = append (received , got )
60+ total += int (got .NumRows ())
6061 }
6162 require .Empty (t , plugin .RecordsDiff (table .ToArrowSchema (), []arrow.Record {record }, received ))
6263 require .NoError (t , readErr )
63- require .Equalf (t , rows , len ( received ) , "got %d row(s), want %d" , len ( received ) , rows )
64+ require .Equalf (t , rows , total , "got %d row(s), want %d" , total , rows )
6465}
6566func TestWriteReadSliced (t * testing.T ) {
6667 const rows = 10
@@ -102,13 +103,22 @@ func TestWriteReadSliced(t *testing.T) {
102103 readErr = cl .Read (byteReader , table , ch )
103104 close (ch )
104105 }()
105- received := make ([]arrow.Record , 0 , rows )
106+ received , total := make ([]arrow.Record , 0 , rows ), 0
106107 for got := range ch {
107108 received = append (received , got )
109+ total += int (got .NumRows ())
108110 }
109111 require .Empty (t , plugin .RecordsDiff (table .ToArrowSchema (), []arrow.Record {record }, received ))
110112 require .NoError (t , readErr )
111- require .Equalf (t , rows , len (received ), "got %d row(s), want %d" , len (received ), rows )
113+ require .Equalf (t , rows , total , "got %d row(s), want %d" , total , rows )
114+ }
115+
116+ func slice (r arrow.Record ) []arrow.Record {
117+ res := make ([]arrow.Record , r .NumRows ())
118+ for i := int64 (0 ); i < r .NumRows (); i ++ {
119+ res [i ] = r .NewSlice (i , i + 1 )
120+ }
121+ return res
112122}
113123
114124func BenchmarkWrite (b * testing.B ) {
0 commit comments