github.com/parquet-go/parquet-go@v0.21.1-0.20240501160520-b3c3a0c3ed6f/filter_test.go (about) 1 package parquet_test 2 3 import ( 4 "testing" 5 6 "github.com/parquet-go/parquet-go" 7 ) 8 9 func TestFilterRowReader(t *testing.T) { 10 rows := []parquet.Row{ 11 {parquet.Int64Value(0)}, 12 {parquet.Int64Value(1)}, 13 {parquet.Int64Value(2)}, 14 {parquet.Int64Value(3)}, 15 {parquet.Int64Value(4)}, 16 } 17 18 want := []parquet.Row{ 19 {parquet.Int64Value(0)}, 20 {parquet.Int64Value(2)}, 21 {parquet.Int64Value(4)}, 22 } 23 24 reader := parquet.FilterRowReader(&bufferedRows{rows: rows}, 25 func(row parquet.Row) bool { 26 return row[0].Int64()%2 == 0 27 }, 28 ) 29 30 writer := &bufferedRows{} 31 _, err := parquet.CopyRows(writer, reader) 32 if err != nil { 33 t.Fatal(err) 34 } 35 36 assertEqualRows(t, want, writer.rows) 37 } 38 39 func TestFilterRowWriter(t *testing.T) { 40 rows := []parquet.Row{ 41 {parquet.Int64Value(0)}, 42 {parquet.Int64Value(1)}, 43 {parquet.Int64Value(2)}, 44 {parquet.Int64Value(3)}, 45 {parquet.Int64Value(4)}, 46 } 47 48 want := []parquet.Row{ 49 {parquet.Int64Value(1)}, 50 {parquet.Int64Value(3)}, 51 } 52 53 buffer := &bufferedRows{} 54 writer := parquet.FilterRowWriter(buffer, 55 func(row parquet.Row) bool { 56 return row[0].Int64()%2 == 1 57 }, 58 ) 59 60 reader := &bufferedRows{rows: rows} 61 _, err := parquet.CopyRows(writer, reader) 62 if err != nil { 63 t.Fatal(err) 64 } 65 66 assertEqualRows(t, want, buffer.rows) 67 }