fathom/internal/store/chunk.go

package store

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"math"
)

// ColType is a column's physical encoding.
type ColType uint8

const (
	ColDict    ColType = iota // uint32 codes into the chunk dictionary
	ColInt64                  // timestamps, counts, attribute offsets
	ColUint8                  // small enums: kind, status, severity
	ColFloat32                // reserved: no column currently uses it
	ColBlobRef                // offset+length into the shard's .blob sidecar
	ColFloat64                // sample rate, where float32's 1e-7 relative
	//                           error would show up as 9.99999985 in a count
)

// BlobRef points at large text in the shard's blob sidecar. Length zero means
// absent, which is the common case: most spans carry no input or output.
type BlobRef struct {
	Off int64
	Len uint32
}

// Column is one physical column. Exactly one of the typed slices is populated,
// selected by Type. A tagged union would be tidier, but Go's version costs an
// interface dispatch per value in the scan's hot loop, which is the one loop in
// this program worth keeping tight.
type Column struct {
	Name string
	Type ColType

	Codes []uint32
	I64   []int64
	U8    []uint8
	F32   []float32
	F64   []float64
	Blobs []BlobRef
}

func (c *Column) len() int {
	switch c.Type {
	case ColDict:
		return len(c.Codes)
	case ColInt64:
		return len(c.I64)
	case ColUint8:
		return len(c.U8)
	case ColFloat32:
		return len(c.F32)
	case ColFloat64:
		return len(c.F64)
	case ColBlobRef:
		return len(c.Blobs)
	}
	return 0
}

// Attribute columns are ordinary Int64 columns holding a row's window into the
// chunk-wide attribute side table. Keeping them as normal columns rather than
// special-casing them means the scanner has one code path, not two.
const (
	colAttrStart = "_attr_start"
	colAttrCount = "_attr_count"
)

// Chunk is a self-contained, independently decodable batch of rows.
//
// Self-contained is the important property. Each chunk carries its own
// dictionary and its own complete column data, so a reader can decode any chunk
// without having seen the ones before it, and a writer can append one without
// revising anything already on disk. That is what makes the shard file
// append-only and what makes a torn write recoverable by discarding a single
// trailing chunk.
type Chunk struct {
	Dict *Dict
	Cols []Column
	Rows int

	// Attribute side table, indexed by the window in _attr_start/_attr_count.
	AttrKeys  []uint32
	AttrVals  []uint32
	AttrTypes []uint8

	byName map[string]int
}

// Col returns a column by name.
func (c *Chunk) Col(name string) (*Column, bool) {
	if c.byName == nil {
		c.byName = make(map[string]int, len(c.Cols))
		for i := range c.Cols {
			c.byName[c.Cols[i].Name] = i
		}
	}
	i, ok := c.byName[name]
	if !ok {
		return nil, false
	}
	return &c.Cols[i], true
}

const (
	chunkMagic   = "FTHC"
	chunkVersion = 1

	// maxColumns and maxRows bound allocation when decoding a header whose
	// length prefixes may be garbage.
	maxColumns = 256
	maxRows    = 1 << 24
)

// Encode writes the chunk body. The caller frames it with a length and CRC;
// see shard.go.
func (c *Chunk) Encode(w io.Writer) error {
	var scratch [binary.MaxVarintLen64]byte
	putU := func(v uint64) error {
		n := binary.PutUvarint(scratch[:], v)
		_, err := w.Write(scratch[:n])
		return err
	}

	if _, err := io.WriteString(w, chunkMagic); err != nil {
		return err
	}
	if err := putU(chunkVersion); err != nil {
		return err
	}
	if err := putU(uint64(c.Rows)); err != nil {
		return err
	}
	if err := c.Dict.encode(w); err != nil {
		return err
	}

	if err := putU(uint64(len(c.Cols))); err != nil {
		return err
	}
	for i := range c.Cols {
		col := &c.Cols[i]
		if got := col.len(); got != c.Rows {
			return fmt.Errorf("column %q has %d values, chunk has %d rows", col.Name, got, c.Rows)
		}
		if err := putU(uint64(len(col.Name))); err != nil {
			return err
		}
		if _, err := io.WriteString(w, col.Name); err != nil {
			return err
		}
		if err := putU(uint64(col.Type)); err != nil {
			return err
		}
		if err := encodeColumnData(w, col, scratch[:]); err != nil {
			return err
		}
	}

	// Attribute side table.
	if err := putU(uint64(len(c.AttrKeys))); err != nil {
		return err
	}
	for i := range c.AttrKeys {
		if err := putU(uint64(c.AttrKeys[i])); err != nil {
			return err
		}
		if err := putU(uint64(c.AttrVals[i])); err != nil {
			return err
		}
		if _, err := w.Write([]byte{c.AttrTypes[i]}); err != nil {
			return err
		}
	}
	return nil
}

func encodeColumnData(w io.Writer, col *Column, scratch []byte) error {
	putU := func(v uint64) error {
		n := binary.PutUvarint(scratch, v)
		_, err := w.Write(scratch[:n])
		return err
	}
	switch col.Type {
	case ColDict:
		for _, v := range col.Codes {
			if err := putU(uint64(v)); err != nil {
				return err
			}
		}
	case ColInt64:
		// Zig-zag: deltas and negatives are both common here (parent offsets,
		// clock-skewed durations), and varint alone would spend ten bytes on
		// every negative value.
		for _, v := range col.I64 {
			if err := putU(uint64(zigzag(v))); err != nil {
				return err
			}
		}
	case ColUint8:
		if _, err := w.Write(col.U8); err != nil {
			return err
		}
	case ColFloat32:
		var b [4]byte
		for _, v := range col.F32 {
			binary.LittleEndian.PutUint32(b[:], f32bits(v))
			if _, err := w.Write(b[:]); err != nil {
				return err
			}
		}
	case ColFloat64:
		var b [8]byte
		for _, v := range col.F64 {
			binary.LittleEndian.PutUint64(b[:], math.Float64bits(v))
			if _, err := w.Write(b[:]); err != nil {
				return err
			}
		}
	case ColBlobRef:
		for _, v := range col.Blobs {
			if err := putU(uint64(v.Off)); err != nil {
				return err
			}
			if err := putU(uint64(v.Len)); err != nil {
				return err
			}
		}
	default:
		return fmt.Errorf("unknown column type %d for column %q", col.Type, col.Name)
	}
	return nil
}

// DecodeChunk reads a chunk body produced by Encode.
func DecodeChunk(b []byte) (*Chunk, error) {
	r := &reader{b: b}

	magic, err := r.fixed(len(chunkMagic))
	if err != nil {
		return nil, err
	}
	if !bytes.Equal(magic, []byte(chunkMagic)) {
		return nil, fmt.Errorf("bad chunk magic %q", magic)
	}
	ver, err := r.uvarint()
	if err != nil {
		return nil, err
	}
	if ver != chunkVersion {
		return nil, fmt.Errorf("chunk version %d, this build understands %d", ver, chunkVersion)
	}
	rows, err := r.uvarint()
	if err != nil {
		return nil, err
	}
	if rows > maxRows {
		return nil, fmt.Errorf("chunk claims %d rows, over the %d cap", rows, maxRows)
	}
	c := &Chunk{Rows: int(rows)}

	if c.Dict, err = decodeDict(r); err != nil {
		return nil, err
	}

	ncols, err := r.uvarint()
	if err != nil {
		return nil, err
	}
	if ncols > maxColumns {
		return nil, fmt.Errorf("chunk claims %d columns, over the %d cap", ncols, maxColumns)
	}
	c.Cols = make([]Column, 0, ncols)
	for i := uint64(0); i < ncols; i++ {
		name, err := r.str()
		if err != nil {
			return nil, err
		}
		t, err := r.uvarint()
		if err != nil {
			return nil, err
		}
		col := Column{Name: name, Type: ColType(t)}
		if err := decodeColumnData(r, &col, int(rows)); err != nil {
			return nil, fmt.Errorf("column %q: %w", name, err)
		}
		c.Cols = append(c.Cols, col)
	}

	nattrs, err := r.uvarint()
	if err != nil {
		return nil, err
	}
	if nattrs > maxRows*maxColumns {
		return nil, fmt.Errorf("chunk claims %d attribute pairs, implausible", nattrs)
	}
	c.AttrKeys = make([]uint32, 0, nattrs)
	c.AttrVals = make([]uint32, 0, nattrs)
	c.AttrTypes = make([]uint8, 0, nattrs)
	for i := uint64(0); i < nattrs; i++ {
		k, err := r.uvarint()
		if err != nil {
			return nil, err
		}
		v, err := r.uvarint()
		if err != nil {
			return nil, err
		}
		t, err := r.byte()
		if err != nil {
			return nil, err
		}
		c.AttrKeys = append(c.AttrKeys, uint32(k))
		c.AttrVals = append(c.AttrVals, uint32(v))
		c.AttrTypes = append(c.AttrTypes, t)
	}
	return c, nil
}

func decodeColumnData(r *reader, col *Column, rows int) error {
	switch col.Type {
	case ColDict:
		col.Codes = make([]uint32, rows)
		for i := range col.Codes {
			v, err := r.uvarint()
			if err != nil {
				return err
			}
			col.Codes[i] = uint32(v)
		}
	case ColInt64:
		col.I64 = make([]int64, rows)
		for i := range col.I64 {
			v, err := r.uvarint()
			if err != nil {
				return err
			}
			col.I64[i] = unzigzag(v)
		}
	case ColUint8:
		b, err := r.fixed(rows)
		if err != nil {
			return err
		}
		col.U8 = make([]uint8, rows)
		copy(col.U8, b)
	case ColFloat32:
		col.F32 = make([]float32, rows)
		for i := range col.F32 {
			b, err := r.fixed(4)
			if err != nil {
				return err
			}
			col.F32[i] = f32frombits(binary.LittleEndian.Uint32(b))
		}
	case ColFloat64:
		col.F64 = make([]float64, rows)
		for i := range col.F64 {
			b, err := r.fixed(8)
			if err != nil {
				return err
			}
			col.F64[i] = math.Float64frombits(binary.LittleEndian.Uint64(b))
		}
	case ColBlobRef:
		col.Blobs = make([]BlobRef, rows)
		for i := range col.Blobs {
			off, err := r.uvarint()
			if err != nil {
				return err
			}
			n, err := r.uvarint()
			if err != nil {
				return err
			}
			col.Blobs[i] = BlobRef{Off: int64(off), Len: uint32(n)}
		}
	default:
		return fmt.Errorf("unknown column type %d", col.Type)
	}
	return nil
}

func zigzag(v int64) uint64        { return uint64(v<<1) ^ uint64(v>>63) }
func unzigzag(v uint64) int64      { return int64(v>>1) ^ -int64(v&1) }
func f32bits(f float32) uint32     { return math.Float32bits(f) }
func f32frombits(u uint32) float32 { return math.Float32frombits(u) }