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

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