github.com/grailbio/bigslice@v0.0.0-20230519005545-30c4c12152ad/reduce_test.go (about)

     1  // Copyright 2018 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache 2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  package bigslice_test
     6  
     7  import (
     8  	"fmt"
     9  	"testing"
    10  
    11  	"github.com/grailbio/bigslice"
    12  	"github.com/grailbio/bigslice/slicetest"
    13  )
    14  
    15  func TestReduce(t *testing.T) {
    16  	const N = 100
    17  	ints := make([]int, N)
    18  	for i := range ints {
    19  		ints[i] = i
    20  	}
    21  	for m := 1; m < 5; m++ {
    22  		slice := bigslice.Const(m, ints)
    23  		slice = bigslice.Map(slice, func(x int) (string, int) {
    24  			return fmt.Sprint(x%3) + "x", x
    25  		})
    26  		slice = bigslice.Reduce(slice, func(x, y int) int { return x + y })
    27  		assertEqual(t, slice, true, []string{"0x", "1x", "2x"}, []int{1683, 1617, 1650})
    28  	}
    29  }
    30  
    31  func TestReducePrefix(t *testing.T) {
    32  	const N = 100
    33  	ints := make([]int, N)
    34  	for i := range ints {
    35  		ints[i] = i
    36  	}
    37  	for m := 1; m < 5; m++ {
    38  		slice := bigslice.Const(m, ints)
    39  		slice = bigslice.Map(slice, func(x int) (string, int, int) {
    40  			return "x", x % 3, x
    41  		})
    42  		slice = bigslice.Prefixed(slice, 2)
    43  		slice = bigslice.Reduce(slice, func(x, y int) int { return x + y })
    44  		assertEqual(t, slice, true, []string{"x", "x", "x"}, []int{0, 1, 2}, []int{1683, 1617, 1650})
    45  	}
    46  }
    47  
    48  func ExampleReduce() {
    49  	slice := bigslice.Const(2,
    50  		[]string{"c", "a", "b", "c", "c", "b", "a", "a", "a", "a", "c"},
    51  		[]int{1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1},
    52  	)
    53  	slice = bigslice.Reduce(slice, func(a, b int) int { return a + b })
    54  	slicetest.Print(slice)
    55  	// Output:
    56  	// a 5
    57  	// b 2
    58  	// c 4
    59  }