go-hep.org/x/hep@v0.38.1/groot/rarrow/rarrow_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 "io" 10 "strings" 11 "testing" 12 13 "git.sr.ht/~sbinet/go-arrow" 14 "go-hep.org/x/hep/groot" 15 "go-hep.org/x/hep/groot/riofs" 16 "go-hep.org/x/hep/groot/rtree" 17 ) 18 19 func TestSchemaFrom(t *testing.T) { 20 for _, tc := range []struct { 21 file string 22 tree string 23 want *arrow.Schema 24 }{ 25 { 26 file: "../testdata/simple.root", 27 tree: "tree", 28 want: arrow.NewSchema([]arrow.Field{ 29 {Name: "one", Type: arrow.PrimitiveTypes.Int32}, 30 {Name: "two", Type: arrow.PrimitiveTypes.Float32}, 31 {Name: "three", Type: arrow.BinaryTypes.String}, 32 }, nil), 33 }, 34 { 35 file: "../testdata/small-flat-tree.root", 36 tree: "tree", 37 want: arrow.NewSchema([]arrow.Field{ 38 {Name: "Int32", Type: arrow.PrimitiveTypes.Int32}, 39 {Name: "Int64", Type: arrow.PrimitiveTypes.Int64}, 40 {Name: "UInt32", Type: arrow.PrimitiveTypes.Uint32}, 41 {Name: "UInt64", Type: arrow.PrimitiveTypes.Uint64}, 42 {Name: "Float32", Type: arrow.PrimitiveTypes.Float32}, 43 {Name: "Float64", Type: arrow.PrimitiveTypes.Float64}, 44 {Name: "Str", Type: arrow.BinaryTypes.String}, 45 {Name: "ArrayInt32", Type: arrow.FixedSizeListOf(10, arrow.PrimitiveTypes.Int32)}, 46 {Name: "ArrayInt64", Type: arrow.FixedSizeListOf(10, arrow.PrimitiveTypes.Int64)}, 47 {Name: "ArrayUInt32", Type: arrow.FixedSizeListOf(10, arrow.PrimitiveTypes.Uint32)}, 48 {Name: "ArrayUInt64", Type: arrow.FixedSizeListOf(10, arrow.PrimitiveTypes.Uint64)}, 49 {Name: "ArrayFloat32", Type: arrow.FixedSizeListOf(10, arrow.PrimitiveTypes.Float32)}, 50 {Name: "ArrayFloat64", Type: arrow.FixedSizeListOf(10, arrow.PrimitiveTypes.Float64)}, 51 {Name: "N", Type: arrow.PrimitiveTypes.Int32}, 52 {Name: "SliceInt32", Type: arrow.ListOf(arrow.PrimitiveTypes.Int32)}, 53 {Name: "SliceInt64", Type: arrow.ListOf(arrow.PrimitiveTypes.Int64)}, 54 {Name: "SliceUInt32", Type: arrow.ListOf(arrow.PrimitiveTypes.Uint32)}, 55 {Name: "SliceUInt64", Type: arrow.ListOf(arrow.PrimitiveTypes.Uint64)}, 56 {Name: "SliceFloat32", Type: arrow.ListOf(arrow.PrimitiveTypes.Float32)}, 57 {Name: "SliceFloat64", Type: arrow.ListOf(arrow.PrimitiveTypes.Float64)}, 58 }, nil), 59 }, 60 { 61 file: "../testdata/small-evnt-tree-fullsplit.root", 62 tree: "tree", 63 want: arrow.NewSchema([]arrow.Field{ 64 {Name: "evt", Type: arrow.StructOf([]arrow.Field{ 65 {Name: "Beg", Type: arrow.BinaryTypes.String}, 66 {Name: "I16", Type: arrow.PrimitiveTypes.Int16}, 67 {Name: "I32", Type: arrow.PrimitiveTypes.Int32}, 68 {Name: "I64", Type: arrow.PrimitiveTypes.Int64}, 69 {Name: "U16", Type: arrow.PrimitiveTypes.Uint16}, 70 {Name: "U32", Type: arrow.PrimitiveTypes.Uint32}, 71 {Name: "U64", Type: arrow.PrimitiveTypes.Uint64}, 72 {Name: "F32", Type: arrow.PrimitiveTypes.Float32}, 73 {Name: "F64", Type: arrow.PrimitiveTypes.Float64}, 74 {Name: "Str", Type: arrow.BinaryTypes.String}, 75 {Name: "P3", Type: arrow.StructOf([]arrow.Field{ 76 {Name: "Px", Type: arrow.PrimitiveTypes.Int32}, 77 {Name: "Py", Type: arrow.PrimitiveTypes.Float64}, 78 {Name: "Pz", Type: arrow.PrimitiveTypes.Int32}, 79 }...)}, 80 {Name: "ArrayI16", Type: arrow.FixedSizeListOf(10, arrow.PrimitiveTypes.Int16)}, 81 {Name: "ArrayI32", Type: arrow.FixedSizeListOf(10, arrow.PrimitiveTypes.Int32)}, 82 {Name: "ArrayI64", Type: arrow.FixedSizeListOf(10, arrow.PrimitiveTypes.Int64)}, 83 {Name: "ArrayU16", Type: arrow.FixedSizeListOf(10, arrow.PrimitiveTypes.Uint16)}, 84 {Name: "ArrayU32", Type: arrow.FixedSizeListOf(10, arrow.PrimitiveTypes.Uint32)}, 85 {Name: "ArrayU64", Type: arrow.FixedSizeListOf(10, arrow.PrimitiveTypes.Uint64)}, 86 {Name: "ArrayF32", Type: arrow.FixedSizeListOf(10, arrow.PrimitiveTypes.Float32)}, 87 {Name: "ArrayF64", Type: arrow.FixedSizeListOf(10, arrow.PrimitiveTypes.Float64)}, 88 {Name: "N", Type: arrow.PrimitiveTypes.Int32}, 89 {Name: "SliceI16", Type: arrow.ListOf(arrow.PrimitiveTypes.Int16)}, 90 {Name: "SliceI32", Type: arrow.ListOf(arrow.PrimitiveTypes.Int32)}, 91 {Name: "SliceI64", Type: arrow.ListOf(arrow.PrimitiveTypes.Int64)}, 92 {Name: "SliceU16", Type: arrow.ListOf(arrow.PrimitiveTypes.Uint16)}, 93 {Name: "SliceU32", Type: arrow.ListOf(arrow.PrimitiveTypes.Uint32)}, 94 {Name: "SliceU64", Type: arrow.ListOf(arrow.PrimitiveTypes.Uint64)}, 95 {Name: "SliceF32", Type: arrow.ListOf(arrow.PrimitiveTypes.Float32)}, 96 {Name: "SliceF64", Type: arrow.ListOf(arrow.PrimitiveTypes.Float64)}, 97 {Name: "StdStr", Type: arrow.BinaryTypes.String}, 98 {Name: "StlVecI16", Type: arrow.ListOf(arrow.PrimitiveTypes.Int16)}, 99 {Name: "StlVecI32", Type: arrow.ListOf(arrow.PrimitiveTypes.Int32)}, 100 {Name: "StlVecI64", Type: arrow.ListOf(arrow.PrimitiveTypes.Int64)}, 101 {Name: "StlVecU16", Type: arrow.ListOf(arrow.PrimitiveTypes.Uint16)}, 102 {Name: "StlVecU32", Type: arrow.ListOf(arrow.PrimitiveTypes.Uint32)}, 103 {Name: "StlVecU64", Type: arrow.ListOf(arrow.PrimitiveTypes.Uint64)}, 104 {Name: "StlVecF32", Type: arrow.ListOf(arrow.PrimitiveTypes.Float32)}, 105 {Name: "StlVecF64", Type: arrow.ListOf(arrow.PrimitiveTypes.Float64)}, 106 {Name: "StlVecStr", Type: arrow.ListOf(arrow.BinaryTypes.String)}, 107 {Name: "End", Type: arrow.BinaryTypes.String}, 108 }...)}, 109 }, nil), 110 }, 111 { 112 file: "../testdata/small-evnt-tree-nosplit.root", 113 tree: "tree", 114 want: arrow.NewSchema([]arrow.Field{ 115 {Name: "evt", Type: arrow.StructOf([]arrow.Field{ 116 {Name: "Beg", Type: arrow.BinaryTypes.String}, 117 {Name: "I16", Type: arrow.PrimitiveTypes.Int16}, 118 {Name: "I32", Type: arrow.PrimitiveTypes.Int32}, 119 {Name: "I64", Type: arrow.PrimitiveTypes.Int64}, 120 {Name: "U16", Type: arrow.PrimitiveTypes.Uint16}, 121 {Name: "U32", Type: arrow.PrimitiveTypes.Uint32}, 122 {Name: "U64", Type: arrow.PrimitiveTypes.Uint64}, 123 {Name: "F32", Type: arrow.PrimitiveTypes.Float32}, 124 {Name: "F64", Type: arrow.PrimitiveTypes.Float64}, 125 {Name: "Str", Type: arrow.BinaryTypes.String}, 126 {Name: "P3", Type: arrow.StructOf([]arrow.Field{ 127 {Name: "Px", Type: arrow.PrimitiveTypes.Int32}, 128 {Name: "Py", Type: arrow.PrimitiveTypes.Float64}, 129 {Name: "Pz", Type: arrow.PrimitiveTypes.Int32}, 130 }...)}, 131 {Name: "ArrayI16", Type: arrow.FixedSizeListOf(10, arrow.PrimitiveTypes.Int16)}, 132 {Name: "ArrayI32", Type: arrow.FixedSizeListOf(10, arrow.PrimitiveTypes.Int32)}, 133 {Name: "ArrayI64", Type: arrow.FixedSizeListOf(10, arrow.PrimitiveTypes.Int64)}, 134 {Name: "ArrayU16", Type: arrow.FixedSizeListOf(10, arrow.PrimitiveTypes.Uint16)}, 135 {Name: "ArrayU32", Type: arrow.FixedSizeListOf(10, arrow.PrimitiveTypes.Uint32)}, 136 {Name: "ArrayU64", Type: arrow.FixedSizeListOf(10, arrow.PrimitiveTypes.Uint64)}, 137 {Name: "ArrayF32", Type: arrow.FixedSizeListOf(10, arrow.PrimitiveTypes.Float32)}, 138 {Name: "ArrayF64", Type: arrow.FixedSizeListOf(10, arrow.PrimitiveTypes.Float64)}, 139 {Name: "N", Type: arrow.PrimitiveTypes.Int32}, 140 {Name: "SliceI16", Type: arrow.ListOf(arrow.PrimitiveTypes.Int16)}, 141 {Name: "SliceI32", Type: arrow.ListOf(arrow.PrimitiveTypes.Int32)}, 142 {Name: "SliceI64", Type: arrow.ListOf(arrow.PrimitiveTypes.Int64)}, 143 {Name: "SliceU16", Type: arrow.ListOf(arrow.PrimitiveTypes.Uint16)}, 144 {Name: "SliceU32", Type: arrow.ListOf(arrow.PrimitiveTypes.Uint32)}, 145 {Name: "SliceU64", Type: arrow.ListOf(arrow.PrimitiveTypes.Uint64)}, 146 {Name: "SliceF32", Type: arrow.ListOf(arrow.PrimitiveTypes.Float32)}, 147 {Name: "SliceF64", Type: arrow.ListOf(arrow.PrimitiveTypes.Float64)}, 148 {Name: "StdStr", Type: arrow.BinaryTypes.String}, 149 {Name: "StlVecI16", Type: arrow.ListOf(arrow.PrimitiveTypes.Int16)}, 150 {Name: "StlVecI32", Type: arrow.ListOf(arrow.PrimitiveTypes.Int32)}, 151 {Name: "StlVecI64", Type: arrow.ListOf(arrow.PrimitiveTypes.Int64)}, 152 {Name: "StlVecU16", Type: arrow.ListOf(arrow.PrimitiveTypes.Uint16)}, 153 {Name: "StlVecU32", Type: arrow.ListOf(arrow.PrimitiveTypes.Uint32)}, 154 {Name: "StlVecU64", Type: arrow.ListOf(arrow.PrimitiveTypes.Uint64)}, 155 {Name: "StlVecF32", Type: arrow.ListOf(arrow.PrimitiveTypes.Float32)}, 156 {Name: "StlVecF64", Type: arrow.ListOf(arrow.PrimitiveTypes.Float64)}, 157 {Name: "StlVecStr", Type: arrow.ListOf(arrow.BinaryTypes.String)}, 158 {Name: "End", Type: arrow.BinaryTypes.String}, 159 }...)}, 160 }, nil), 161 }, 162 { 163 file: "../testdata/root_numpy_struct.root", 164 tree: "test", 165 want: arrow.NewSchema([]arrow.Field{ 166 {Name: "branch1", Type: arrow.StructOf([]arrow.Field{ 167 {Name: "intleaf", Type: arrow.PrimitiveTypes.Int32}, 168 {Name: "floatleaf", Type: arrow.PrimitiveTypes.Float32}, 169 }...)}, 170 {Name: "branch2", Type: arrow.StructOf([]arrow.Field{ 171 {Name: "intleaf", Type: arrow.PrimitiveTypes.Int32}, 172 {Name: "floatleaf", Type: arrow.PrimitiveTypes.Float32}, 173 }...)}, 174 }, nil), 175 }, 176 } { 177 t.Run(tc.file, func(t *testing.T) { 178 f, err := groot.Open(tc.file) 179 if err != nil { 180 t.Fatal(err) 181 } 182 defer f.Close() 183 184 o, err := riofs.Dir(f).Get(tc.tree) 185 if err != nil { 186 t.Fatal(err) 187 } 188 189 tree := o.(rtree.Tree) 190 191 got := SchemaFrom(tree) 192 193 if !got.Equal(tc.want) { 194 t.Fatalf("invalid schema.\ngot:\n%s\nwant:\n%s\n", displaySchema(got), displaySchema(tc.want)) 195 } 196 }) 197 } 198 } 199 200 func displaySchema(s *arrow.Schema) string { 201 o := new(strings.Builder) 202 fmt.Fprintf(o, "%*.sfields: %d\n", 2, "", len(s.Fields())) 203 for _, f := range s.Fields() { 204 displayField(o, f, 4) 205 } 206 if meta := s.Metadata(); meta.Len() > 0 { 207 fmt.Fprintf(o, "metadata: %v\n", meta) 208 } 209 return o.String() 210 } 211 212 func displayField(o io.Writer, field arrow.Field, inc int) { 213 nullable := "" 214 if field.Nullable { 215 nullable = ", nullable" 216 } 217 fmt.Fprintf(o, "%*.s- %s: type=%v%v\n", inc, "", field.Name, field.Type, nullable) 218 if field.HasMetadata() { 219 fmt.Fprintf(o, "%*.smetadata: %v\n", inc, "", field.Metadata) 220 } 221 }