github.com/grailbio/bigslice@v0.0.0-20230519005545-30c4c12152ad/exec/combiner_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 exec
     6  
     7  import (
     8  	"bytes"
     9  	"context"
    10  	"reflect"
    11  	"sort"
    12  	"testing"
    13  
    14  	fuzz "github.com/google/gofuzz"
    15  	"github.com/grailbio/bigslice/frame"
    16  	"github.com/grailbio/bigslice/slicefunc"
    17  	"github.com/grailbio/bigslice/sliceio"
    18  	"github.com/grailbio/bigslice/slicetype"
    19  )
    20  
    21  var typeOfInt = reflect.TypeOf(0)
    22  
    23  func deepEqual(f, g frame.Frame) bool {
    24  	if f.NumOut() != g.NumOut() {
    25  		return false
    26  	}
    27  	for i := 0; i < f.NumOut(); i++ {
    28  		if f.Out(i) != g.Out(i) {
    29  			return false
    30  		}
    31  		if !reflect.DeepEqual(f.Interface(i), g.Interface(i)) {
    32  			return false
    33  		}
    34  	}
    35  	return true
    36  }
    37  
    38  func TestCombiningFrame(t *testing.T) {
    39  	typ := slicetype.New(typeOfString, typeOfInt)
    40  	fn, ok := slicefunc.Of(func(n, m int) int { return n + m })
    41  	if !ok {
    42  		t.Fatal("unexpected bad func")
    43  	}
    44  	f := makeCombiningFrame(typ, fn, 2, 1)
    45  	if f == nil {
    46  		t.Fatal("nil frame")
    47  	}
    48  	f.Combine(frame.Slices(
    49  		[]string{"a", "b", "a", "a", "a"},
    50  		[]int{1, 2, 10, 20, 30},
    51  	))
    52  	f.Combine(frame.Slices(
    53  		[]string{"x", "a", "a"},
    54  		[]int{100, 0, 0},
    55  	))
    56  	if got, want := f.Len(), 3; got != want {
    57  		t.Errorf("got %v, want %v", got, want)
    58  	}
    59  	g := f.Compact()
    60  	sort.Sort(g)
    61  	if got, want := g.Interface(0).([]string), ([]string{"a", "b", "x"}); !reflect.DeepEqual(got, want) {
    62  		t.Errorf("got %v, want %v", got, want)
    63  	}
    64  	if got, want := g.Interface(1).([]int), ([]int{61, 2, 100}); !reflect.DeepEqual(got, want) {
    65  		t.Errorf("got %v, want %v", got, want)
    66  	}
    67  }
    68  
    69  func TestCombiningFrameManyKeys(t *testing.T) {
    70  	const N = 100000
    71  	typ := slicetype.New(typeOfString, typeOfInt)
    72  	fn, ok := slicefunc.Of(func(n, m int) int { return n + m })
    73  	if !ok {
    74  		t.Fatal("unexpected bad func")
    75  	}
    76  	f := makeCombiningFrame(typ, fn, 2, 1)
    77  	if f == nil {
    78  		t.Fatal("nil frame")
    79  	}
    80  	fz := fuzz.New()
    81  	fz.NilChance(0)
    82  	total := make(map[string]int)
    83  	for i := 0; i < N; i++ {
    84  		var elems map[string]int
    85  		fz.Fuzz(&elems)
    86  		// add some common ones too
    87  		elems["a"] = 123
    88  		elems["b"] = 333
    89  		var (
    90  			ks []string
    91  			vs []int
    92  		)
    93  		for k, v := range elems {
    94  			ks = append(ks, k)
    95  			vs = append(vs, v)
    96  			total[k] += v
    97  		}
    98  		f.Combine(frame.Slices(ks, vs))
    99  
   100  	}
   101  	c := f.Compact()
   102  	sort.Sort(c)
   103  	keys := make([]string, 0, len(total))
   104  	for k := range total {
   105  		keys = append(keys, k)
   106  	}
   107  	sort.Strings(keys)
   108  	if got, want := c.Len(), len(total); got != want {
   109  		t.Fatalf("got %v, want %v", got, want)
   110  	}
   111  	for i, key := range keys {
   112  		if got, want := c.Index(0, i).String(), key; got != want {
   113  			t.Fatalf("index %d: got %v, want %v", i, got, want)
   114  		}
   115  		if got, want := c.Index(1, i).Int(), int64(total[key]); got != want {
   116  			t.Errorf("index %d: got %v, want %v", i, got, want)
   117  		}
   118  	}
   119  }
   120  
   121  func TestCombiner(t *testing.T) {
   122  	const N = 100
   123  	typ := slicetype.New(typeOfString, typeOfInt)
   124  	fn, ok := slicefunc.Of(func(n, m int) int { return n + m })
   125  	if !ok {
   126  		t.Fatal("unexpected bad func")
   127  	}
   128  	// Set a small target value to ensure spilling.
   129  	c, err := newCombiner(typ, "test", fn, 2)
   130  	if err != nil {
   131  		t.Fatal(err)
   132  	}
   133  	ctx := context.Background()
   134  	f := frame.Slices(
   135  		[]string{"a", "a", "b", "c", "d"},
   136  		[]int{0, 1, 2, 3, 4},
   137  	)
   138  	for i := 0; i < N; i++ {
   139  		if err = c.Combine(ctx, f); err != nil {
   140  			t.Fatal(err)
   141  		}
   142  	}
   143  	var b bytes.Buffer
   144  	n, err := c.WriteTo(ctx, sliceio.NewEncodingWriter(&b))
   145  	if err != nil {
   146  		t.Fatal(err)
   147  	}
   148  	if got, want := n, int64(4); got != want {
   149  		t.Errorf("got %v, want %v", got, want)
   150  	}
   151  	r := sliceio.NewDecodingReader(&b)
   152  	g := frame.Make(f, int(n), int(n))
   153  	m, err := sliceio.ReadFull(ctx, r, g)
   154  	if err != nil {
   155  		t.Fatal(err)
   156  	}
   157  	if got, want := m, 4; got != want {
   158  		t.Errorf("got %v, want %v", got, want)
   159  	}
   160  	if got, want := g, frame.Slices([]string{"a", "b", "c", "d"}, []int{N, 2 * N, 3 * N, 4 * N}); !deepEqual(got, want) {
   161  		t.Errorf("got %v, want %v", got.TabString(), want.TabString())
   162  	}
   163  }