github.com/segmentio/parquet-go@v0.0.0-20230712180008-5d42db8f0d47/transform_test.go (about)

     1  package parquet_test
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/segmentio/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  }