diff --git a/groot/riofs/file.go b/groot/riofs/file.go index 9547547e..204b643b 100644 --- a/groot/riofs/file.go +++ b/groot/riofs/file.go @@ -574,7 +574,7 @@ func (f *File) writeStreamerInfo() error { return fmt.Errorf("riofs: could not write StreamerInfo list: %w", err) } - key, err = newKeyFromBuf(&f.dir, "StreamerInfo", sinfos.Title(), sinfos.Class(), 1, buf.Bytes(), f) + key, err = newKeyFromBuf(&f.dir, "StreamerInfo", sinfos.Title(), sinfos.Class(), 1, buf.Bytes(), f, nil) if err != nil { return fmt.Errorf("riofs: could not create StreamerInfo key: %w", err) } diff --git a/groot/riofs/key.go b/groot/riofs/key.go index 048629b7..80741da4 100644 --- a/groot/riofs/key.go +++ b/groot/riofs/key.go @@ -18,6 +18,27 @@ import ( "go-hep.org/x/hep/groot/rvers" ) +// KeyOption configures how a Key may be created. +type KeyOption func(cfg *keyCfg) error + +// WithKeyCompression configures a Key to use the provided compression scheme. +func WithKeyCompression(compression int32) KeyOption { + return func(cfg *keyCfg) error { + cfg.compr = rcompress.SettingsFrom(compression) + return nil + } +} + +type keyCfg struct { + compr rcompress.Settings +} + +func newKeyCfg() *keyCfg { + return &keyCfg{ + compr: rcompress.Settings{Alg: rcompress.Inherit}, + } +} + // noKeyError is the error returned when a riofs.Key could not be found. type noKeyError struct { key string @@ -122,12 +143,12 @@ func newKey(dir *tdirectoryFile, name, title, class string, objlen int32, f *Fil // NewKey creates a new key from the provided serialized object buffer. // NewKey puts the key and its payload at the end of the provided file f. // Depending on the file configuration, NewKey may compress the provided object buffer. -func NewKey(dir Directory, name, title, class string, cycle int16, obj []byte, f *File) (Key, error) { +func NewKey(dir Directory, name, title, class string, cycle int16, obj []byte, f *File, kopts ...KeyOption) (Key, error) { var d *tdirectoryFile if dir != nil { d = dir.(*tdirectoryFile) } - return newKeyFromBuf(d, name, title, class, cycle, obj, f) + return newKeyFromBuf(d, name, title, class, cycle, obj, f, kopts) } func newKeyFrom(dir *tdirectoryFile, name, title, class string, obj root.Object, f *File) (Key, error) { @@ -185,12 +206,20 @@ func newKeyFrom(dir *tdirectoryFile, name, title, class string, obj root.Object, return k, nil } -func newKeyFromBuf(dir *tdirectoryFile, name, title, class string, cycle int16, buf []byte, f *File) (Key, error) { +func newKeyFromBuf(dir *tdirectoryFile, name, title, class string, cycle int16, buf []byte, f *File, kopts []KeyOption) (Key, error) { var err error if dir == nil { dir = &f.dir } + kcfg := newKeyCfg() + for _, opt := range kopts { + err = opt(kcfg) + if err != nil { + return Key{}, fmt.Errorf("riofs: could not setup Key option: %w", err) + } + } + keylen := keylenFor(name, title, class, dir, f.end) objlen := int32(len(buf)) k := Key{ @@ -212,7 +241,12 @@ func newKeyFromBuf(dir *tdirectoryFile, name, title, class string, cycle int16, k.rvers += 1000 } - k.buf, err = rcompress.Compress(nil, buf, k.f.compression) + compress := k.f.compression + if kcfg.compr.Alg != rcompress.Inherit { + compress = kcfg.compr.Compression() + } + + k.buf, err = rcompress.Compress(nil, buf, compress) if err != nil { return k, fmt.Errorf("riofs: could not compress object %s for key %q: %w", class, name, err) } diff --git a/groot/rtree/basket.go b/groot/rtree/basket.go index 58d21028..f5061ded 100644 --- a/groot/rtree/basket.go +++ b/groot/rtree/basket.go @@ -306,7 +306,7 @@ func (b *Basket) grow(n int) { b.offsets = append(b.offsets, make([]int32, delta)...) } -func (b *Basket) writeFile(f *riofs.File) (totBytes int64, zipBytes int64, err error) { +func (b *Basket) writeFile(f *riofs.File, compr int32) (totBytes int64, zipBytes int64, err error) { header := b.header b.header = true defer func() { @@ -332,7 +332,7 @@ func (b *Basket) writeFile(f *riofs.File) (totBytes int64, zipBytes int64, err e b.wbuf.WriteArrayI32(b.offsets[:b.nevbuf]) b.wbuf.WriteI32(0) } - b.key, err = riofs.NewKey(nil, b.key.Name(), b.key.Title(), b.Class(), int16(b.key.Cycle()), b.wbuf.Bytes(), f) + b.key, err = riofs.NewKey(nil, b.key.Name(), b.key.Title(), b.Class(), int16(b.key.Cycle()), b.wbuf.Bytes(), f, riofs.WithKeyCompression(compr)) if err != nil { return 0, 0, fmt.Errorf("rtree: could not create basket-key: %w", err) } diff --git a/groot/rtree/branch.go b/groot/rtree/branch.go index c47fa128..d3c7569b 100644 --- a/groot/rtree/branch.go +++ b/groot/rtree/branch.go @@ -749,7 +749,7 @@ func (b *tbranch) flush() error { } f := b.tree.getFile() - totBytes, zipBytes, err := b.ctx.bk.writeFile(f) + totBytes, zipBytes, err := b.ctx.bk.writeFile(f, int32(b.compress)) if err != nil { return fmt.Errorf("could not marshal basket[%d] (branch=%q): %w", b.writeBasket, b.Name(), err) }