go-hep.org/x/hep@v0.38.1/groot/rnpy/arrow_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 rnpy
     6  
     7  import (
     8  	"os"
     9  	"path/filepath"
    10  	"testing"
    11  
    12  	"codeberg.org/sbinet/npyio"
    13  )
    14  
    15  func TestRecord(t *testing.T) {
    16  	for _, tc := range []struct {
    17  		name   string
    18  		forder bool // Fortran order
    19  		want   any
    20  	}{
    21  		{
    22  			name: "float64_4x3x2-c-order",
    23  			want: [4][3][2]float64{
    24  				{{10, 11}, {12, 13}, {14, 15}},
    25  				{{16, 17}, {18, 19}, {20, 21}},
    26  				{{22, 23}, {24, 25}, {26, 27}},
    27  				{{28, 29}, {30, 31}, {32, 33}},
    28  			},
    29  		},
    30  		{
    31  			name:   "float64_4x3x2-f-order",
    32  			forder: true,
    33  			want: [2][3][4]float64{
    34  				{{10, 11, 12, 13}, {14, 15, 16, 17}, {18, 19, 20, 21}},
    35  				{{22, 23, 24, 25}, {26, 27, 28, 29}, {30, 31, 32, 33}},
    36  			},
    37  		},
    38  	} {
    39  		t.Run(tc.name, func(t *testing.T) {
    40  			tmp, err := os.MkdirTemp("", "npy2root-")
    41  			if err != nil {
    42  				t.Fatalf("%+v", err)
    43  			}
    44  			defer os.RemoveAll(tmp)
    45  
    46  			fname := filepath.Join(tmp, "data.npy")
    47  			src, err := os.Create(fname)
    48  			if err != nil {
    49  				t.Fatalf("could not create NumPy data file: %+v", err)
    50  			}
    51  			defer src.Close()
    52  
    53  			err = npyio.Write(src, tc.want)
    54  			if err != nil {
    55  				t.Fatalf("could not save NumPy data file: %+v", err)
    56  			}
    57  
    58  			err = src.Close()
    59  			if err != nil {
    60  				t.Fatalf("could not close NumPy data file: %+v", err)
    61  			}
    62  
    63  			src, err = os.Open(fname)
    64  			if err != nil {
    65  				t.Fatalf("could not reopen NumPy data file: %+v", err)
    66  			}
    67  			defer src.Close()
    68  
    69  			npy, err := npyio.NewReader(src)
    70  			if err != nil {
    71  				t.Fatalf("could not create numpy file reader: %+v", err)
    72  			}
    73  
    74  			npy.Header.Descr.Fortran = tc.forder
    75  
    76  			rec := NewRecord(npy)
    77  			defer rec.Release()
    78  
    79  			rec.Retain()
    80  			rec.Release()
    81  
    82  			if got, want := rec.NumRows(), int64(4); got != want {
    83  				t.Fatalf("invalid number of rows: got=%d, want=%d", got, want)
    84  			}
    85  
    86  			if got, want := rec.NumCols(), int64(1); got != want {
    87  				t.Fatalf("invalid number of cols: got=%d, want=%d", got, want)
    88  			}
    89  
    90  			if got, want := rec.ColumnName(0), "numpy"; got != want {
    91  				t.Fatalf("invalid column name: got=%q, want=%q", got, want)
    92  			}
    93  		})
    94  	}
    95  }