github.com/grailbio/base@v0.0.11/recordio/example_indexing_test.go (about)

     1  package recordio_test
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/gob"
     6  	"fmt"
     7  	"io"
     8  
     9  	"github.com/grailbio/base/recordio"
    10  )
    11  
    12  type recordioIndex map[string]recordio.ItemLocation
    13  
    14  func doWriteWithIndex(out io.Writer) {
    15  	index := make(recordioIndex)
    16  	wr := recordio.NewWriter(out, recordio.WriterOpts{
    17  		Marshal: func(scratch []byte, v interface{}) ([]byte, error) { return []byte(v.(string)), nil },
    18  		Index: func(loc recordio.ItemLocation, val interface{}) error {
    19  			index[val.(string)] = loc
    20  			return nil
    21  		},
    22  	})
    23  
    24  	// To store a trailer block, AddHeader(recordio.KeyTrailer, true) must be
    25  	// called beforehand.
    26  	wr.AddHeader(recordio.KeyTrailer, true)
    27  	wr.Append("Item0")
    28  	wr.Append("Item1")
    29  	wr.Append("Item2")
    30  	wr.Flush()
    31  	// Wait for the index callbacks to run.
    32  	wr.Wait()
    33  
    34  	// Write the index in the trailer.
    35  	indexBuf := &bytes.Buffer{}
    36  	encoder := gob.NewEncoder(indexBuf)
    37  	if err := encoder.Encode(index); err != nil {
    38  		panic(err)
    39  	}
    40  	wr.SetTrailer(indexBuf.Bytes())
    41  	if err := wr.Finish(); err != nil {
    42  		panic(err)
    43  	}
    44  }
    45  
    46  func doReadWithIndex(in io.ReadSeeker) {
    47  	r := recordio.NewScanner(in, recordio.ScannerOpts{
    48  		Unmarshal: func(data []byte) (interface{}, error) { return string(data), nil },
    49  	})
    50  	// Read the trailer, parse it into the recordioIndex.
    51  	decoder := gob.NewDecoder(bytes.NewReader(r.Trailer()))
    52  	index := make(recordioIndex)
    53  	if err := decoder.Decode(&index); err != nil {
    54  		panic(err)
    55  	}
    56  	// Try reading individual items.
    57  	r.Seek(index["Item1"])
    58  	for r.Scan() {
    59  		fmt.Printf("Item: %s\n", r.Get().(string))
    60  	}
    61  	r.Seek(index["Item0"])
    62  	for r.Scan() {
    63  		fmt.Printf("Item: %s\n", r.Get().(string))
    64  	}
    65  	if err := r.Err(); err != nil {
    66  		panic(err)
    67  	}
    68  }
    69  
    70  func Example_indexing() {
    71  	buf := &bytes.Buffer{}
    72  	doWriteWithIndex(buf)
    73  	doReadWithIndex(bytes.NewReader(buf.Bytes()))
    74  	// Output:
    75  	// Item: Item1
    76  	// Item: Item2
    77  	// Item: Item0
    78  	// Item: Item1
    79  	// Item: Item2
    80  }