groot/{internal/rcompress,riofs,rtree}: expose API to customize key-c…

…ompression Signed-off-by: Sebastien Binet <[email protected]>
go-hep · Oct 22, 2024 · aab80b1 · aab80b1
1 parent 7c6853a
commit aab80b1
Show file tree

Hide file tree

Showing 6 changed files with 66 additions and 11 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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)
 	}