go-hep.org/x/hep@v0.38.1/cmd/arrow2root/main.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  // arrow2root converts the content of an ARROW file to a ROOT TTree.
     6  package main // import "go-hep.org/x/hep/cmd/arrow2root"
     7  
     8  import (
     9  	"flag"
    10  	"fmt"
    11  	"log"
    12  	"os"
    13  
    14  	"git.sr.ht/~sbinet/go-arrow/arrio"
    15  	"git.sr.ht/~sbinet/go-arrow/ipc"
    16  	"git.sr.ht/~sbinet/go-arrow/memory"
    17  	"go-hep.org/x/hep/groot"
    18  	"go-hep.org/x/hep/groot/rarrow"
    19  	"go-hep.org/x/hep/groot/rtree"
    20  )
    21  
    22  func main() {
    23  	log.SetPrefix("arrow2root: ")
    24  	log.SetFlags(0)
    25  
    26  	oname := flag.String("o", "output.root", "path to output ROOT file name")
    27  	tname := flag.String("t", "tree", "name of the output tree")
    28  
    29  	flag.Parse()
    30  
    31  	if flag.NArg() != 1 {
    32  		flag.Usage()
    33  		log.Fatalf("missing input ARROW filename argument")
    34  	}
    35  	fname := flag.Arg(0)
    36  
    37  	err := process(*oname, *tname, fname)
    38  	if err != nil {
    39  		log.Fatalf("%+v", err)
    40  	}
    41  }
    42  
    43  func process(oname, tname, fname string) error {
    44  	f, err := os.Open(fname)
    45  	if err != nil {
    46  		return fmt.Errorf("could not open ARROW file %q: %w", fname, err)
    47  	}
    48  	defer f.Close()
    49  
    50  	mem := memory.NewGoAllocator()
    51  	r, err := ipc.NewFileReader(f, ipc.WithAllocator(mem))
    52  	if err != nil {
    53  		return fmt.Errorf("could not create ARROW IPC reader from %q: %w", fname, err)
    54  	}
    55  	defer r.Close()
    56  
    57  	o, err := groot.Create(oname)
    58  	if err != nil {
    59  		return fmt.Errorf("could not create output ROOT file %q: %w", oname, err)
    60  	}
    61  	defer o.Close()
    62  
    63  	tree, err := rarrow.NewFlatTreeWriter(o, tname, r.Schema(), rtree.WithTitle(tname))
    64  	if err != nil {
    65  		return fmt.Errorf("could not create output ROOT tree %q: %w", tname, err)
    66  	}
    67  
    68  	_, err = arrio.Copy(tree, r)
    69  	if err != nil {
    70  		return fmt.Errorf("could not convert ARROW file to ROOT tree: %w", err)
    71  	}
    72  
    73  	err = tree.Close()
    74  	if err != nil {
    75  		return fmt.Errorf("could not close ROOT tree writer: %w", err)
    76  	}
    77  
    78  	err = o.Close()
    79  	if err != nil {
    80  		return fmt.Errorf("could not close output ROOT file %q: %w", oname, err)
    81  	}
    82  
    83  	return nil
    84  }