go-hep.org/x/hep@v0.38.1/groot/rarrow/table.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  	"sync/atomic"
    10  
    11  	"git.sr.ht/~sbinet/go-arrow"
    12  	"git.sr.ht/~sbinet/go-arrow/array"
    13  	"git.sr.ht/~sbinet/go-arrow/memory"
    14  	"go-hep.org/x/hep/groot/rtree"
    15  )
    16  
    17  // NewTable creates a new in-memory Arrow Table from the provided ROOT Tree.
    18  func NewTable(t rtree.Tree, opts ...Option) array.Table {
    19  	cfg := newConfig(opts)
    20  
    21  	tbl := &rootTable{
    22  		mem:    cfg.mem,
    23  		tree:   t,
    24  		refs:   1,
    25  		schema: SchemaFrom(t),
    26  		nrows:  t.Entries(),
    27  		ncols:  int64(len(t.Branches())),
    28  		cols:   make([]*array.Column, len(t.Branches())),
    29  	}
    30  
    31  	tbl.init()
    32  
    33  	return tbl
    34  }
    35  
    36  type rootTable struct {
    37  	mem  memory.Allocator
    38  	tree rtree.Tree
    39  
    40  	refs   int64
    41  	schema *arrow.Schema
    42  	nrows  int64
    43  	ncols  int64
    44  
    45  	cols []*array.Column
    46  }
    47  
    48  func (tbl *rootTable) Schema() *arrow.Schema      { return tbl.schema }
    49  func (tbl *rootTable) NumRows() int64             { return tbl.nrows }
    50  func (tbl *rootTable) NumCols() int64             { return tbl.ncols }
    51  func (tbl *rootTable) Column(i int) *array.Column { return tbl.cols[i] }
    52  
    53  // Retain increases the reference count by 1.
    54  // Retain may be called simultaneously from multiple goroutines.
    55  func (tbl *rootTable) Retain() {
    56  	atomic.AddInt64(&tbl.refs, 1)
    57  }
    58  
    59  // Release decreases the reference count by 1.
    60  // When the reference count goes to zero, the memory is freed.
    61  // Release may be called simultaneously from multiple goroutines.
    62  func (tbl *rootTable) Release() {
    63  	if atomic.LoadInt64(&tbl.refs) <= 0 {
    64  		panic("groot/rarrow: too many releases")
    65  	}
    66  
    67  	if atomic.AddInt64(&tbl.refs, -1) == 0 {
    68  		for i := range tbl.cols {
    69  			tbl.cols[i].Release()
    70  		}
    71  		tbl.cols = nil
    72  	}
    73  }
    74  
    75  func (tbl *rootTable) init() {
    76  	// FIXME(sbinet): infer clusters sizes
    77  	// FIXME(sbinet): lazily populate rootTable
    78  
    79  	var (
    80  		rvars  = rtree.NewReadVars(tbl.tree)
    81  		r, err = rtree.NewReader(tbl.tree, rvars)
    82  	)
    83  	if err != nil {
    84  		panic(fmt.Errorf("could not create reader from read-vars %#v: %+v", rvars, err))
    85  	}
    86  	defer r.Close()
    87  
    88  	arrs := make([]array.Interface, tbl.ncols)
    89  	blds := make([]array.Builder, tbl.ncols)
    90  	for i, field := range tbl.schema.Fields() {
    91  		blds[i] = builderFrom(tbl.mem, field.Type, tbl.nrows)
    92  		defer blds[i].Release()
    93  	}
    94  
    95  	err = r.Read(func(ctx rtree.RCtx) error {
    96  		for i, field := range tbl.schema.Fields() {
    97  			appendData(blds[i], rvars[i], field.Type)
    98  		}
    99  		return nil
   100  	})
   101  	if err != nil {
   102  		panic(fmt.Errorf("could not read tree: %+v", err))
   103  	}
   104  
   105  	for i, bldr := range blds {
   106  		arrs[i] = bldr.NewArray()
   107  		defer arrs[i].Release()
   108  	}
   109  
   110  	tbl.cols = make([]*array.Column, tbl.ncols)
   111  	for i, arr := range arrs {
   112  		field := tbl.schema.Field(i)
   113  		if !arrow.TypeEqual(field.Type, arr.DataType()) {
   114  			panic(fmt.Errorf("field[%d][%s]: type=%v|%v array=%v", i, field.Name, field.Type, arr.DataType(), arr))
   115  		}
   116  		chunked := array.NewChunked(field.Type, []array.Interface{arr})
   117  		defer chunked.Release()
   118  		tbl.cols[i] = array.NewColumn(field, chunked)
   119  	}
   120  }
   121  
   122  var (
   123  	_ array.Table = (*rootTable)(nil)
   124  )