From 5400b84b8e94a686e57560d87c0f76d90950c7df Mon Sep 17 00:00:00 2001 From: Alfonso Subiotto Marques Date: Wed, 29 Nov 2023 15:12:49 +0100 Subject: [PATCH] snapshot: release arrow readers after use This is good practice, and we're possibly retaining unnecessary memory here. --- snapshot.go | 28 +++++++++++++++++++--------- 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/snapshot.go b/snapshot.go index 3a4cbe3b0..f5f94af2f 100644 --- a/snapshot.go +++ b/snapshot.go @@ -634,17 +634,27 @@ func loadSnapshot(ctx context.Context, db *DB, r io.ReaderAt, size int64) ([]byt } resultParts = append(resultParts, parts.NewPart(partMeta.Tx, serBuf, partOptions)) case snapshotpb.Part_ENCODING_ARROW: - arrowReader, err := ipc.NewReader(bytes.NewReader(partBytes)) - if err != nil { - return err - } - - record, err := arrowReader.Read() - if err != nil { + if err := func() error { + arrowReader, err := ipc.NewReader(bytes.NewReader(partBytes)) + if err != nil { + return err + } + defer arrowReader.Release() + + record, err := arrowReader.Read() + if err != nil { + return err + } + + record.Retain() + resultParts = append( + resultParts, + parts.NewArrowPart(partMeta.Tx, record, int(util.TotalRecordSize(record)), table.schema, partOptions), + ) + return nil + }(); err != nil { return err } - - resultParts = append(resultParts, parts.NewArrowPart(partMeta.Tx, record, int(util.TotalRecordSize(record)), table.schema, partOptions)) default: return fmt.Errorf("unknown part encoding: %s", partMeta.Encoding) }