groot/{rarrow,rcmd,rdict,riofs}: properly handle ndim-arrays

Fixes #807.
go-hep · Oct 6, 2020 · 0be35e7 · 0be35e7
1 parent 4d49a78
commit 0be35e7
Show file tree

Hide file tree

Showing 14 changed files with 442 additions and 9 deletions.
diff --git a/groot/rarrow/rarrow.go b/groot/rarrow/rarrow.go
@@ -51,6 +51,7 @@ func dataTypeFromLeaf(leaf rtree.Leaf) arrow.DataType {
 	var (
 		unsigned = leaf.IsUnsigned()
 		kind     = leaf.Kind()
+		typ      = leaf.Type()
 		dt       arrow.DataType
 	)
 
@@ -93,10 +94,10 @@ func dataTypeFromLeaf(leaf rtree.Leaf) arrow.DataType {
 		dt = arrow.BinaryTypes.String
 
 	case reflect.Struct:
-		dt = dataTypeFromGo(leaf.Type())
+		dt = dataTypeFromGo(typ)
 
 	case reflect.Slice:
-		dt = dataTypeFromGo(leaf.Type())
+		dt = dataTypeFromGo(typ)
 
 	default:
 		panic(fmt.Errorf("not implemented %#v (kind=%v)", leaf, kind))
@@ -118,7 +119,15 @@ func dataTypeFromLeaf(leaf rtree.Leaf) arrow.DataType {
 				panic(fmt.Errorf("groot/rtree: invalid number of dimensions (%d)", dims))
 			}
 		default:
-			dt = arrow.FixedSizeListOf(int32(leaf.Len()), dt)
+			shape := leaf.Shape()
+			switch leaf.(type) {
+			case *rtree.LeafF16, *rtree.LeafD32:
+				// workaround for https://sft.its.cern.ch/jira/browse/ROOT-10149
+				shape = []int{leaf.Len()}
+			}
+			for i := range shape {
+				dt = arrow.FixedSizeListOf(int32(shape[len(shape)-1-i]), dt)
+			}
 		}
 	}
 

diff --git a/groot/rarrow/reader_test.go b/groot/rarrow/reader_test.go
@@ -43,6 +43,15 @@ func TestRecord(t *testing.T) {
 			tree: "tree",
 			want: "testdata/small-evnt-tree-nosplit.root.txt",
 		},
+		{
+			// n-dim arrays
+			// FIXME(sbinet): arrays of Float16_t and Double32_t are flatten.
+			// This is because of:
+			// https://sft.its.cern.ch/jira/browse/ROOT-10149
+			file: "../testdata/ndim.root",
+			tree: "tree",
+			want: "testdata/ndim.root.txt",
+		},
 	} {
 		t.Run(tc.file, func(t *testing.T) {
 			f, err := groot.Open(tc.file)

diff --git a/groot/rarrow/testdata/ndim.root.txt b/groot/rarrow/testdata/ndim.root.txt
@@ -0,0 +1,14 @@
+file: ../testdata/ndim.root
+rec[0][ArrBs]: [[[[[true false true false true] [true false true false true] [true false true false true] [true false true false true]] [[true false true false true] [true false true false true] [true false true false true] [true false true false true]] [[true false true false true] [true false true false true] [true false true false true] [true false true false true]]] [[[true false true false true] [true false true false true] [true false true false true] [true false true false true]] [[true false true false true] [true false true false true] [true false true false true] [true false true false true]] [[true false true false true] [true false true false true] [true false true false true] [true false true false true]]]] [[[[true false true false true] [true false true false true] [true false true false true] [true false true false true]] [[true false true false true] [true false true false true] [true false true false true] [true false true false true]] [[true false true false true] [true false true false true] [true false true false true] [true false true false true]]] [[[true false true false true] [true false true false true] [true false true false true] [true false true false true]] [[true false true false true] [true false true false true] [true false true false true] [true false true false true]] [[true false true false true] [true false true false true] [true false true false true] [true false true false true]]]]]
+rec[0][ArrI8]: [[[[[0 -1 -2 -3 -4] [-5 -6 -7 -8 -9] [-10 -11 -12 -13 -14] [-15 -16 -17 -18 -19]] [[-20 -21 -22 -23 -24] [-25 -26 -27 -28 -29] [-30 -31 -32 -33 -34] [-35 -36 -37 -38 -39]] [[-40 -41 -42 -43 -44] [-45 -46 -47 -48 -49] [-50 -51 -52 -53 -54] [-55 -56 -57 -58 -59]]] [[[-60 -61 -62 -63 -64] [-65 -66 -67 -68 -69] [-70 -71 -72 -73 -74] [-75 -76 -77 -78 -79]] [[-80 -81 -82 -83 -84] [-85 -86 -87 -88 -89] [-90 -91 -92 -93 -94] [-95 -96 -97 -98 -99]] [[-100 -101 -102 -103 -104] [-105 -106 -107 -108 -109] [-110 -111 -112 -113 -114] [-115 -116 -117 -118 -119]]]] [[[[-1 -2 -3 -4 -5] [-6 -7 -8 -9 -10] [-11 -12 -13 -14 -15] [-16 -17 -18 -19 -20]] [[-21 -22 -23 -24 -25] [-26 -27 -28 -29 -30] [-31 -32 -33 -34 -35] [-36 -37 -38 -39 -40]] [[-41 -42 -43 -44 -45] [-46 -47 -48 -49 -50] [-51 -52 -53 -54 -55] [-56 -57 -58 -59 -60]]] [[[-61 -62 -63 -64 -65] [-66 -67 -68 -69 -70] [-71 -72 -73 -74 -75] [-76 -77 -78 -79 -80]] [[-81 -82 -83 -84 -85] [-86 -87 -88 -89 -90] [-91 -92 -93 -94 -95] [-96 -97 -98 -99 -100]] [[-101 -102 -103 -104 -105] [-106 -107 -108 -109 -110] [-111 -112 -113 -114 -115] [-116 -117 -118 -119 -120]]]]]
+rec[0][ArrI16]: [[[[[0 -1 -2 -3 -4] [-5 -6 -7 -8 -9] [-10 -11 -12 -13 -14] [-15 -16 -17 -18 -19]] [[-20 -21 -22 -23 -24] [-25 -26 -27 -28 -29] [-30 -31 -32 -33 -34] [-35 -36 -37 -38 -39]] [[-40 -41 -42 -43 -44] [-45 -46 -47 -48 -49] [-50 -51 -52 -53 -54] [-55 -56 -57 -58 -59]]] [[[-60 -61 -62 -63 -64] [-65 -66 -67 -68 -69] [-70 -71 -72 -73 -74] [-75 -76 -77 -78 -79]] [[-80 -81 -82 -83 -84] [-85 -86 -87 -88 -89] [-90 -91 -92 -93 -94] [-95 -96 -97 -98 -99]] [[-100 -101 -102 -103 -104] [-105 -106 -107 -108 -109] [-110 -111 -112 -113 -114] [-115 -116 -117 -118 -119]]]] [[[[-1 -2 -3 -4 -5] [-6 -7 -8 -9 -10] [-11 -12 -13 -14 -15] [-16 -17 -18 -19 -20]] [[-21 -22 -23 -24 -25] [-26 -27 -28 -29 -30] [-31 -32 -33 -34 -35] [-36 -37 -38 -39 -40]] [[-41 -42 -43 -44 -45] [-46 -47 -48 -49 -50] [-51 -52 -53 -54 -55] [-56 -57 -58 -59 -60]]] [[[-61 -62 -63 -64 -65] [-66 -67 -68 -69 -70] [-71 -72 -73 -74 -75] [-76 -77 -78 -79 -80]] [[-81 -82 -83 -84 -85] [-86 -87 -88 -89 -90] [-91 -92 -93 -94 -95] [-96 -97 -98 -99 -100]] [[-101 -102 -103 -104 -105] [-106 -107 -108 -109 -110] [-111 -112 -113 -114 -115] [-116 -117 -118 -119 -120]]]]]
+rec[0][ArrI32]: [[[[[0 -1 -2 -3 -4] [-5 -6 -7 -8 -9] [-10 -11 -12 -13 -14] [-15 -16 -17 -18 -19]] [[-20 -21 -22 -23 -24] [-25 -26 -27 -28 -29] [-30 -31 -32 -33 -34] [-35 -36 -37 -38 -39]] [[-40 -41 -42 -43 -44] [-45 -46 -47 -48 -49] [-50 -51 -52 -53 -54] [-55 -56 -57 -58 -59]]] [[[-60 -61 -62 -63 -64] [-65 -66 -67 -68 -69] [-70 -71 -72 -73 -74] [-75 -76 -77 -78 -79]] [[-80 -81 -82 -83 -84] [-85 -86 -87 -88 -89] [-90 -91 -92 -93 -94] [-95 -96 -97 -98 -99]] [[-100 -101 -102 -103 -104] [-105 -106 -107 -108 -109] [-110 -111 -112 -113 -114] [-115 -116 -117 -118 -119]]]] [[[[-1 -2 -3 -4 -5] [-6 -7 -8 -9 -10] [-11 -12 -13 -14 -15] [-16 -17 -18 -19 -20]] [[-21 -22 -23 -24 -25] [-26 -27 -28 -29 -30] [-31 -32 -33 -34 -35] [-36 -37 -38 -39 -40]] [[-41 -42 -43 -44 -45] [-46 -47 -48 -49 -50] [-51 -52 -53 -54 -55] [-56 -57 -58 -59 -60]]] [[[-61 -62 -63 -64 -65] [-66 -67 -68 -69 -70] [-71 -72 -73 -74 -75] [-76 -77 -78 -79 -80]] [[-81 -82 -83 -84 -85] [-86 -87 -88 -89 -90] [-91 -92 -93 -94 -95] [-96 -97 -98 -99 -100]] [[-101 -102 -103 -104 -105] [-106 -107 -108 -109 -110] [-111 -112 -113 -114 -115] [-116 -117 -118 -119 -120]]]]]
+rec[0][ArrI64]: [[[[[0 -1 -2 -3 -4] [-5 -6 -7 -8 -9] [-10 -11 -12 -13 -14] [-15 -16 -17 -18 -19]] [[-20 -21 -22 -23 -24] [-25 -26 -27 -28 -29] [-30 -31 -32 -33 -34] [-35 -36 -37 -38 -39]] [[-40 -41 -42 -43 -44] [-45 -46 -47 -48 -49] [-50 -51 -52 -53 -54] [-55 -56 -57 -58 -59]]] [[[-60 -61 -62 -63 -64] [-65 -66 -67 -68 -69] [-70 -71 -72 -73 -74] [-75 -76 -77 -78 -79]] [[-80 -81 -82 -83 -84] [-85 -86 -87 -88 -89] [-90 -91 -92 -93 -94] [-95 -96 -97 -98 -99]] [[-100 -101 -102 -103 -104] [-105 -106 -107 -108 -109] [-110 -111 -112 -113 -114] [-115 -116 -117 -118 -119]]]] [[[[-1 -2 -3 -4 -5] [-6 -7 -8 -9 -10] [-11 -12 -13 -14 -15] [-16 -17 -18 -19 -20]] [[-21 -22 -23 -24 -25] [-26 -27 -28 -29 -30] [-31 -32 -33 -34 -35] [-36 -37 -38 -39 -40]] [[-41 -42 -43 -44 -45] [-46 -47 -48 -49 -50] [-51 -52 -53 -54 -55] [-56 -57 -58 -59 -60]]] [[[-61 -62 -63 -64 -65] [-66 -67 -68 -69 -70] [-71 -72 -73 -74 -75] [-76 -77 -78 -79 -80]] [[-81 -82 -83 -84 -85] [-86 -87 -88 -89 -90] [-91 -92 -93 -94 -95] [-96 -97 -98 -99 -100]] [[-101 -102 -103 -104 -105] [-106 -107 -108 -109 -110] [-111 -112 -113 -114 -115] [-116 -117 -118 -119 -120]]]]]
+rec[0][ArrU8]: [[[[[0 1 2 3 4] [5 6 7 8 9] [10 11 12 13 14] [15 16 17 18 19]] [[20 21 22 23 24] [25 26 27 28 29] [30 31 32 33 34] [35 36 37 38 39]] [[40 41 42 43 44] [45 46 47 48 49] [50 51 52 53 54] [55 56 57 58 59]]] [[[60 61 62 63 64] [65 66 67 68 69] [70 71 72 73 74] [75 76 77 78 79]] [[80 81 82 83 84] [85 86 87 88 89] [90 91 92 93 94] [95 96 97 98 99]] [[100 101 102 103 104] [105 106 107 108 109] [110 111 112 113 114] [115 116 117 118 119]]]] [[[[1 2 3 4 5] [6 7 8 9 10] [11 12 13 14 15] [16 17 18 19 20]] [[21 22 23 24 25] [26 27 28 29 30] [31 32 33 34 35] [36 37 38 39 40]] [[41 42 43 44 45] [46 47 48 49 50] [51 52 53 54 55] [56 57 58 59 60]]] [[[61 62 63 64 65] [66 67 68 69 70] [71 72 73 74 75] [76 77 78 79 80]] [[81 82 83 84 85] [86 87 88 89 90] [91 92 93 94 95] [96 97 98 99 100]] [[101 102 103 104 105] [106 107 108 109 110] [111 112 113 114 115] [116 117 118 119 120]]]]]
+rec[0][ArrU16]: [[[[[0 1 2 3 4] [5 6 7 8 9] [10 11 12 13 14] [15 16 17 18 19]] [[20 21 22 23 24] [25 26 27 28 29] [30 31 32 33 34] [35 36 37 38 39]] [[40 41 42 43 44] [45 46 47 48 49] [50 51 52 53 54] [55 56 57 58 59]]] [[[60 61 62 63 64] [65 66 67 68 69] [70 71 72 73 74] [75 76 77 78 79]] [[80 81 82 83 84] [85 86 87 88 89] [90 91 92 93 94] [95 96 97 98 99]] [[100 101 102 103 104] [105 106 107 108 109] [110 111 112 113 114] [115 116 117 118 119]]]] [[[[1 2 3 4 5] [6 7 8 9 10] [11 12 13 14 15] [16 17 18 19 20]] [[21 22 23 24 25] [26 27 28 29 30] [31 32 33 34 35] [36 37 38 39 40]] [[41 42 43 44 45] [46 47 48 49 50] [51 52 53 54 55] [56 57 58 59 60]]] [[[61 62 63 64 65] [66 67 68 69 70] [71 72 73 74 75] [76 77 78 79 80]] [[81 82 83 84 85] [86 87 88 89 90] [91 92 93 94 95] [96 97 98 99 100]] [[101 102 103 104 105] [106 107 108 109 110] [111 112 113 114 115] [116 117 118 119 120]]]]]
+rec[0][ArrU32]: [[[[[0 1 2 3 4] [5 6 7 8 9] [10 11 12 13 14] [15 16 17 18 19]] [[20 21 22 23 24] [25 26 27 28 29] [30 31 32 33 34] [35 36 37 38 39]] [[40 41 42 43 44] [45 46 47 48 49] [50 51 52 53 54] [55 56 57 58 59]]] [[[60 61 62 63 64] [65 66 67 68 69] [70 71 72 73 74] [75 76 77 78 79]] [[80 81 82 83 84] [85 86 87 88 89] [90 91 92 93 94] [95 96 97 98 99]] [[100 101 102 103 104] [105 106 107 108 109] [110 111 112 113 114] [115 116 117 118 119]]]] [[[[1 2 3 4 5] [6 7 8 9 10] [11 12 13 14 15] [16 17 18 19 20]] [[21 22 23 24 25] [26 27 28 29 30] [31 32 33 34 35] [36 37 38 39 40]] [[41 42 43 44 45] [46 47 48 49 50] [51 52 53 54 55] [56 57 58 59 60]]] [[[61 62 63 64 65] [66 67 68 69 70] [71 72 73 74 75] [76 77 78 79 80]] [[81 82 83 84 85] [86 87 88 89 90] [91 92 93 94 95] [96 97 98 99 100]] [[101 102 103 104 105] [106 107 108 109 110] [111 112 113 114 115] [116 117 118 119 120]]]]]
+rec[0][ArrU64]: [[[[[0 1 2 3 4] [5 6 7 8 9] [10 11 12 13 14] [15 16 17 18 19]] [[20 21 22 23 24] [25 26 27 28 29] [30 31 32 33 34] [35 36 37 38 39]] [[40 41 42 43 44] [45 46 47 48 49] [50 51 52 53 54] [55 56 57 58 59]]] [[[60 61 62 63 64] [65 66 67 68 69] [70 71 72 73 74] [75 76 77 78 79]] [[80 81 82 83 84] [85 86 87 88 89] [90 91 92 93 94] [95 96 97 98 99]] [[100 101 102 103 104] [105 106 107 108 109] [110 111 112 113 114] [115 116 117 118 119]]]] [[[[1 2 3 4 5] [6 7 8 9 10] [11 12 13 14 15] [16 17 18 19 20]] [[21 22 23 24 25] [26 27 28 29 30] [31 32 33 34 35] [36 37 38 39 40]] [[41 42 43 44 45] [46 47 48 49 50] [51 52 53 54 55] [56 57 58 59 60]]] [[[61 62 63 64 65] [66 67 68 69 70] [71 72 73 74 75] [76 77 78 79 80]] [[81 82 83 84 85] [86 87 88 89 90] [91 92 93 94 95] [96 97 98 99 100]] [[101 102 103 104 105] [106 107 108 109 110] [111 112 113 114 115] [116 117 118 119 120]]]]]
+rec[0][ArrF32]: [[[[[0 1 2 3 4] [5 6 7 8 9] [10 11 12 13 14] [15 16 17 18 19]] [[20 21 22 23 24] [25 26 27 28 29] [30 31 32 33 34] [35 36 37 38 39]] [[40 41 42 43 44] [45 46 47 48 49] [50 51 52 53 54] [55 56 57 58 59]]] [[[60 61 62 63 64] [65 66 67 68 69] [70 71 72 73 74] [75 76 77 78 79]] [[80 81 82 83 84] [85 86 87 88 89] [90 91 92 93 94] [95 96 97 98 99]] [[100 101 102 103 104] [105 106 107 108 109] [110 111 112 113 114] [115 116 117 118 119]]]] [[[[1 2 3 4 5] [6 7 8 9 10] [11 12 13 14 15] [16 17 18 19 20]] [[21 22 23 24 25] [26 27 28 29 30] [31 32 33 34 35] [36 37 38 39 40]] [[41 42 43 44 45] [46 47 48 49 50] [51 52 53 54 55] [56 57 58 59 60]]] [[[61 62 63 64 65] [66 67 68 69 70] [71 72 73 74 75] [76 77 78 79 80]] [[81 82 83 84 85] [86 87 88 89 90] [91 92 93 94 95] [96 97 98 99 100]] [[101 102 103 104 105] [106 107 108 109 110] [111 112 113 114 115] [116 117 118 119 120]]]]]
+rec[0][ArrF64]: [[[[[0 1 2 3 4] [5 6 7 8 9] [10 11 12 13 14] [15 16 17 18 19]] [[20 21 22 23 24] [25 26 27 28 29] [30 31 32 33 34] [35 36 37 38 39]] [[40 41 42 43 44] [45 46 47 48 49] [50 51 52 53 54] [55 56 57 58 59]]] [[[60 61 62 63 64] [65 66 67 68 69] [70 71 72 73 74] [75 76 77 78 79]] [[80 81 82 83 84] [85 86 87 88 89] [90 91 92 93 94] [95 96 97 98 99]] [[100 101 102 103 104] [105 106 107 108 109] [110 111 112 113 114] [115 116 117 118 119]]]] [[[[1 2 3 4 5] [6 7 8 9 10] [11 12 13 14 15] [16 17 18 19 20]] [[21 22 23 24 25] [26 27 28 29 30] [31 32 33 34 35] [36 37 38 39 40]] [[41 42 43 44 45] [46 47 48 49 50] [51 52 53 54 55] [56 57 58 59 60]]] [[[61 62 63 64 65] [66 67 68 69 70] [71 72 73 74 75] [76 77 78 79 80]] [[81 82 83 84 85] [86 87 88 89 90] [91 92 93 94 95] [96 97 98 99 100]] [[101 102 103 104 105] [106 107 108 109 110] [111 112 113 114 115] [116 117 118 119 120]]]]]
+rec[0][ArrD16]: [[0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119] [1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120]]
+rec[0][ArrD32]: [[0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119] [1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120]]
diff --git a/groot/rcmd/dump_test.go b/groot/rcmd/dump_test.go
@@ -206,6 +206,14 @@ func TestDump(t *testing.T) {
 			name: "../testdata/std-map-split1.root",
 			want: loadRef("testdata/std-map-split1.root.txt"),
 		},
+		{
+			// n-dim arrays
+			// FIXME(sbinet): arrays of Float16_t and Double32_t are flatten.
+			// This is because of:
+			// https://sft.its.cern.ch/jira/browse/ROOT-10149
+			name: "../testdata/ndim.root",
+			want: loadRef("testdata/ndim.root.txt"),
+		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			got := new(strings.Builder)