go-hep.org/x/hep@v0.38.1/groot/rarrow/table_test.go (about)

     1  // Copyright ©2019 The go-hep Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package rarrow // import "go-hep.org/x/hep/groot/rarrow"
     6  
     7  import (
     8  	"fmt"
     9  	"os"
    10  	"strings"
    11  	"testing"
    12  
    13  	"git.sr.ht/~sbinet/go-arrow/array"
    14  	"git.sr.ht/~sbinet/go-arrow/memory"
    15  	"go-hep.org/x/hep/groot"
    16  	"go-hep.org/x/hep/groot/riofs"
    17  	"go-hep.org/x/hep/groot/rtree"
    18  )
    19  
    20  func TestTable(t *testing.T) {
    21  	for _, tc := range []struct {
    22  		file string
    23  		tree string
    24  		want string
    25  	}{
    26  		{
    27  			file: "../testdata/simple.root",
    28  			tree: "tree",
    29  			want: "testdata/simple.root.txt",
    30  		},
    31  		{
    32  			file: "../testdata/small-flat-tree.root",
    33  			tree: "tree",
    34  			want: "testdata/small-flat-tree.root.txt",
    35  		},
    36  		{
    37  			file: "../testdata/small-evnt-tree-fullsplit.root",
    38  			tree: "tree",
    39  			want: "testdata/small-evnt-tree-fullsplit.root.txt",
    40  		},
    41  		{
    42  			file: "../testdata/small-evnt-tree-nosplit.root",
    43  			tree: "tree",
    44  			want: "testdata/small-evnt-tree-nosplit.root.txt",
    45  		},
    46  	} {
    47  		t.Run(tc.file, func(t *testing.T) {
    48  			f, err := groot.Open(tc.file)
    49  			if err != nil {
    50  				t.Fatal(err)
    51  			}
    52  			defer f.Close()
    53  
    54  			o, err := riofs.Dir(f).Get(tc.tree)
    55  			if err != nil {
    56  				t.Fatal(err)
    57  			}
    58  
    59  			mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
    60  			defer mem.AssertSize(t, 0)
    61  
    62  			tree := o.(rtree.Tree)
    63  			tbl := NewTable(tree, WithAllocator(mem))
    64  			defer tbl.Release()
    65  
    66  			tbl.Retain()
    67  			tbl.Release()
    68  
    69  			tr := array.NewTableReader(tbl, -1)
    70  			defer tr.Release()
    71  
    72  			recs := 0
    73  			out := new(strings.Builder)
    74  			fmt.Fprintf(out, "file: %s\n", tc.file)
    75  			for tr.Next() {
    76  				rec := tr.Record()
    77  				for i, col := range rec.Columns() {
    78  					fmt.Fprintf(out, "rec[%d][%s]: %v\n", recs, rec.Schema().Field(i).Name, col)
    79  				}
    80  				recs++
    81  			}
    82  
    83  			want, err := os.ReadFile(tc.want)
    84  			if err != nil {
    85  				t.Fatal(err)
    86  			}
    87  
    88  			if got, want := out.String(), string(want); got != want {
    89  				t.Fatalf("invalid table\ngot:\n%s\nwant:\n%s\n", got, want)
    90  			}
    91  		})
    92  	}
    93  }