github.com/parquet-go/parquet-go@v0.21.1-0.20240501160520-b3c3a0c3ed6f/transform_test.go (about) 1 package parquet_test 2 3 import ( 4 "testing" 5 6 "github.com/parquet-go/parquet-go" 7 ) 8 9 func TestTransformRowReader(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), parquet.Int64Value(0).Level(0, 0, 1)}, 20 {parquet.Int64Value(1), parquet.Int64Value(2).Level(0, 0, 1)}, 21 {parquet.Int64Value(2), parquet.Int64Value(4).Level(0, 0, 1)}, 22 {parquet.Int64Value(3), parquet.Int64Value(6).Level(0, 0, 1)}, 23 {parquet.Int64Value(4), parquet.Int64Value(8).Level(0, 0, 1)}, 24 } 25 26 reader := parquet.TransformRowReader(&bufferedRows{rows: rows}, 27 func(dst, src parquet.Row) (parquet.Row, error) { 28 dst = append(dst, src[0]) 29 dst = append(dst, parquet.Int64Value(2*src[0].Int64()).Level(0, 0, 1)) 30 return dst, nil 31 }, 32 ) 33 34 writer := &bufferedRows{} 35 _, err := parquet.CopyRows(writer, reader) 36 if err != nil { 37 t.Fatal(err) 38 } 39 40 assertEqualRows(t, want, writer.rows) 41 } 42 43 func TestTransformRowWriter(t *testing.T) { 44 rows := []parquet.Row{ 45 {parquet.Int64Value(0)}, 46 {parquet.Int64Value(1)}, 47 {parquet.Int64Value(2)}, 48 {parquet.Int64Value(3)}, 49 {parquet.Int64Value(4)}, 50 } 51 52 want := []parquet.Row{ 53 {parquet.Int64Value(1)}, 54 {parquet.Int64Value(3)}, 55 } 56 57 buffer := &bufferedRows{} 58 writer := parquet.TransformRowWriter(buffer, 59 func(dst, src parquet.Row) (parquet.Row, error) { 60 if (src[0].Int64() % 2) != 0 { 61 dst = append(dst, src[0]) 62 } 63 return dst, nil 64 }, 65 ) 66 67 reader := &bufferedRows{rows: rows} 68 _, err := parquet.CopyRows(writer, reader) 69 if err != nil { 70 t.Fatal(err) 71 } 72 73 assertEqualRows(t, want, buffer.rows) 74 }