From aab80b1d6eb7f5fb6af3fdd75373800030197c77 Mon Sep 17 00:00:00 2001 From: Sebastien Binet Date: Tue, 22 Oct 2024 14:51:43 +0200 Subject: [PATCH] groot/{internal/rcompress,riofs,rtree}: expose API to customize key-compression Signed-off-by: Sebastien Binet --- groot/internal/rcompress/rcompress.go | 8 ++++ groot/riofs/dir.go | 2 +- groot/riofs/file.go | 2 +- groot/riofs/key.go | 59 ++++++++++++++++++++++++--- groot/rtree/basket.go | 4 +- groot/rtree/branch.go | 2 +- 6 files changed, 66 insertions(+), 11 deletions(-) diff --git a/groot/internal/rcompress/rcompress.go b/groot/internal/rcompress/rcompress.go index e9b710ab2..75f26eca6 100644 --- a/groot/internal/rcompress/rcompress.go +++ b/groot/internal/rcompress/rcompress.go @@ -55,6 +55,14 @@ type Settings struct { Lvl int } +// SettingsFrom create a Settings value from the provided compression +// configuration (compression algorithm and compression level), using +// ROOT's encoding. +func SettingsFrom(compr int32) Settings { + alg, lvl := rootCompressAlgLvl(compr) + return Settings{Alg: alg, Lvl: lvl} +} + // DefaultSettings is the default compression algorithm and level used // in ROOT files and trees. var DefaultSettings = Settings{Alg: ZLIB, Lvl: flate.BestSpeed} diff --git a/groot/riofs/dir.go b/groot/riofs/dir.go index 5d07c8645..ce22b937d 100644 --- a/groot/riofs/dir.go +++ b/groot/riofs/dir.go @@ -464,7 +464,7 @@ func (dir *tdirectoryFile) Put(name string, obj root.Object) error { dir.addStreamer(si) } - key, err := newKeyFrom(dir, name, title, rdict.GoName2Cxx(typename), obj, dir.file) + key, err := newKeyFrom(dir, name, title, rdict.GoName2Cxx(typename), obj, dir.file, nil) // FIXME(sbinet): wire in key-opt ? if err != nil { return fmt.Errorf("riofs: could not create key %q for object %T: %w", name, obj, err) } diff --git a/groot/riofs/file.go b/groot/riofs/file.go index 9547547e8..204b643ba 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 048629b72..9babd6f6a 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,20 +143,28 @@ 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) { +func newKeyFrom(dir *tdirectoryFile, name, title, class string, obj root.Object, 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) buf := rbytes.NewWBuffer(nil, nil, uint32(keylen), dir.file) @@ -171,7 +200,12 @@ func newKeyFrom(dir *tdirectoryFile, name, title, class string, obj root.Object, k.rvers += 1000 } - k.buf, err = rcompress.Compress(nil, buf.Bytes(), k.f.compression) + compress := k.f.compression + if kcfg.compr.Alg != rcompress.Inherit { + compress = kcfg.compr.Compression() + } + + k.buf, err = rcompress.Compress(nil, buf.Bytes(), compress) if err != nil { return k, fmt.Errorf("riofs: could not compress object %T for key %q: %w", obj, name, err) } @@ -185,12 +219,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 +254,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 58d21028d..f5061dedc 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 c47fa1280..d3c7569b8 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) }