partition.go

package partition

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"io"
	"log"

	"github.com/go-sif/sif"
	pb "github.com/go-sif/sif/internal/rpc"
	iutil "github.com/go-sif/sif/internal/util"
	uuid "github.com/gofrs/uuid"
)

const defaultCapacity = 2

// partitionImpl is Sif's internal implementation of Partition
type partitionImpl struct {
	compressor          sif.Compressor
	schema              sif.Schema
	internalData        *partitionData
	colIsDirty          map[string]bool
	activeColData       map[string][]byte        // uncompressed fixed-width column data
	activeVarLengthData map[string][]interface{} // uncompressed variable-length column data which has been deserialized by the column's deserializer
}

// createPartitionImpl initializes a new Partition, with an initial row capacity, a maximum row capacity, and a particular Schema
func createPartitionImpl(maxRows int, initialCapacity int, schema sif.Schema) *partitionImpl {
	id, err := uuid.NewV4()
	if err != nil {
		log.Fatalf("failed to generate UUID for Partition: %v", err)
	}
	if initialCapacity > maxRows {
		initialCapacity = maxRows
	}
	if initialCapacity < 0 {
		panic("initial capacity for partition cannot be negative")
	}
	// initialize underlying column data
	colData := make(map[string]*columnData)
	colIsDirty := make(map[string]bool)
	activeColData := make(map[string][]byte)
	activeVarLengthData := make(map[string][]interface{})
	schema.ForEachColumn(func(name string, col sif.Column) error {
		isVarLength := !sif.IsFixedWidth(col.Type())
		colIsDirty[name] = true // initially, there is no serialized, compressed data for this column. everything starts as active data
		colData[name] = &columnData{
			IsVariableLength: isVarLength,
			Meta:             make([]byte, initialCapacity),
			Data:             nil, // initially, there is no serialized, compressed data for this column. everything starts as active data
		}
		if isVarLength {
			activeVarLengthData[name] = make([]interface{}, initialCapacity)
		} else {
			fixedCol := col.Type().(sif.FixedWidthColumnType)
			activeColData[name] = make([]byte, initialCapacity*fixedCol.Size())
		}
		return nil
	})
	internalData := &partitionData{
		Id:       id.String(),
		MaxRows:  int32(maxRows),
		NumRows:  0,
		Capacity: int32(initialCapacity),
		Keys:     nil,
		ColData:  colData,
	}
	return &partitionImpl{
		internalData:        internalData,
		colIsDirty:          colIsDirty,
		activeColData:       activeColData,
		activeVarLengthData: activeVarLengthData,
		schema:              schema,
		compressor:          nil,
	}
}

// CreatePartition creates a new Partition containing an empty byte array and a schema
func CreatePartition(maxRows int, initialCapacity int, schema sif.Schema) sif.Partition {
	return createPartitionImpl(maxRows, initialCapacity, schema)
}

// sync writes any Partition modifications to the underlying compressed data,
// returning true iff there were any changes to write
func (p *partitionImpl) sync() (bool, error) {
	// returns error if compressor has not been set
	if p.compressor == nil {
		return true, fmt.Errorf("cannot sync Partition modifications for partition %s because no Compressor has been set", p.internalData.Id)
	}
	anyChanges := false
	for _, colName := range p.schema.ColumnNames() {
		if p.colIsDirty[colName] {
			changed, err := p.serializeColumn(colName)
			if err != nil {
				return anyChanges, err
			}
			anyChanges = anyChanges || changed
		}
	}
	return anyChanges, nil
}

func (p *partitionImpl) columnIsDeserialized(colName string) bool {
	_, isFixedDeserialized := p.activeColData[colName]
	_, isVarDeserialized := p.activeVarLengthData[colName]
	return isFixedDeserialized || isVarDeserialized
}

func (p *partitionImpl) columnIsDirty(colName string) bool {
	isDirty, isKnown := p.colIsDirty[colName]
	return isKnown && isDirty
}

func (p *partitionImpl) deserializeColumn(colName string) error {
	// no-op if already deserialized
	if p.columnIsDeserialized(colName) {
		return nil
	}
	// returns error if decompressor has not been set
	if p.compressor == nil {
		return fmt.Errorf("cannot deserialize column data for column %s because no Compressor has been set", colName)
	}
	col, err := p.schema.GetColumn(colName)
	if err != nil {
		return err
	}
	colData := p.internalData.ColData[colName]
	compressedColData := colData.Data
	// decompress, if there is data to decompress
	var decompressedColData *bytes.Buffer
	if compressedColData != nil {
		decompressedColData = new(bytes.Buffer)
		err = p.compressor.Decompress(bytes.NewBuffer(compressedColData), decompressedColData)
		if err != nil {
			return err
		}
	}
	// Copy decompressed data into active buffers
	if colData.IsVariableLength {
		// make room for active variable-length data
		p.activeVarLengthData[colName] = make([]interface{}, p.internalData.Capacity)
		// if there's no rows, we're done
		if p.internalData.NumRows == 0 {
			return nil
		}
		// if there's no decompressed data we're done
		if decompressedColData == nil {
			return nil
		}
		// decode data from internalData proto
		var decodedSerializedRowData [][]byte
		d := gob.NewDecoder(decompressedColData)
		err := d.Decode(&decodedSerializedRowData)
		if err != nil {
			return err
		}
		if sif.IsFixedWidth(col.Type()) {
			return fmt.Errorf("serialized column data for %s is variable-length, but schema indicates column is fixed-length", colName)
		}
		dest := p.activeVarLengthData[colName]
		for i, sRowVal := range decodedSerializedRowData {
			// we don't need to copy deserialized data past our current partition size
			if i >= p.GetNumRows() {
				break
			}
			if sRowVal != nil {
				dest[i], err = col.Type().Deserialize(sRowVal)
				if err != nil {
					return err
				}
			}
		}
	} else {
		fixedType := col.Type().(sif.FixedWidthColumnType)
		// make room for active fixed-width data
		p.activeColData[colName] = make([]byte, int(p.internalData.Capacity)*fixedType.Size())
		// if there's no rows, we're done
		if p.internalData.NumRows == 0 {
			return nil
		}
		// if there's no decompressed data we're done
		if decompressedColData == nil {
			return nil
		}
		decompressedBytes := decompressedColData.Bytes()
		if len(decompressedBytes) > len(p.activeColData[colName]) {
			return fmt.Errorf("cannot deserialize column data. Size of deserialized column data is larger than partition capacity")
		}
		// we don't need to copy deserialized data past our current partition size
		copy(p.activeColData[colName], decompressedBytes[0:fixedType.Size()*p.GetNumRows()])
	}
	// done, data is ready to be used
	return nil
}

func (p *partitionImpl) serializeColumn(colName string) (bool, error) {
	// no-op if never deserialized
	if !p.columnIsDeserialized(colName) {
		return false, nil
	} else if _, ok := p.internalData.ColData[colName]; ok && !p.columnIsDirty(colName) {
		// nearly a no-op if valid serialized data still exists (discards any deserialized data)
		delete(p.activeColData, colName)
		delete(p.activeVarLengthData, colName)
		return false, nil
	}
	// returns error if compressor has not been set
	if p.compressor == nil {
		return true, fmt.Errorf("cannot serialize column data for column %s because no Compressor has been set", colName)
	}
	if p.internalData.NumRows > 0 {
		toCompress := new(bytes.Buffer)
		colData := p.internalData.ColData[colName]
		col, err := p.schema.GetColumn(colName)
		if err != nil {
			return true, err
		}
		if colData.IsVariableLength {
			// if the data is variable length, we first have to
			// serialize each row value using the Column's Serialize
			// method, then Gob encode the resulting [][]byte into a []byte.
			// Finally, we compress and store in the proto.
			if sif.IsFixedWidth(col.Type()) {
				return true, fmt.Errorf("serialized column data for %s is variable-length, but schema indicates column is fixed-length", colName)
			}
			activeData := p.activeVarLengthData[colName]
			serializedActiveData := make([][]byte, p.internalData.NumRows)
			for i := 0; i < int(p.internalData.NumRows); i++ {
				rowVal := activeData[i]
				if rowVal != nil {
					serializedActiveData[i], err = col.Type().(sif.VariableWidthColumnType).Serialize(rowVal)
					if err != nil {
						return true, err
					}
				}
			}
			e := gob.NewEncoder(toCompress)
			err = e.Encode(serializedActiveData)
			if err != nil {
				return true, err
			}
		} else {
			if !sif.IsFixedWidth(col.Type()) {
				return true, fmt.Errorf("serialized column data for %s is fixed-length, but schema indicates column is variable-length", colName)
			}
			toCompress.Write(p.activeColData[colName][0 : p.internalData.NumRows*int32(col.Type().(sif.FixedWidthColumnType).Size())])
		}
		// compress data
		compressed := new(bytes.Buffer)
		err = p.compressor.Compress(toCompress, compressed)
		if err != nil {
			return true, err
		}
		// store in proto
		colData.Data = compressed.Bytes()
	}
	// column is no longer dirty
	p.colIsDirty[colName] = false
	// discard deserialized data
	delete(p.activeColData, colName)
	delete(p.activeVarLengthData, colName)
	return true, nil
}

// setColumnDirty indicates that a column has been written to and needs to be reserialized via sync()
func (p *partitionImpl) setColumnDirty(colName string) {
	p.colIsDirty[colName] = true
	// remove internal, serialized data for column because it's no longer useful
	p.internalData.ColData[colName].Data = nil
}

// ID retrieves the ID of this Partition
func (p *partitionImpl) ID() string {
	return p.internalData.Id
}

// GetMaxRows retrieves the maximum number of rows in this Partition
func (p *partitionImpl) GetMaxRows() int {
	return int(p.internalData.MaxRows)
}

// GetNumRows retrieves the number of rows in this Partition
func (p *partitionImpl) GetNumRows() int {
	return int(p.internalData.NumRows)
}

// getRowInternal retrieves a specific row from this Partition, without allocation
func (p *partitionImpl) getRow(row *rowImpl, rowNum int) sif.Row {
	row.rowNum = rowNum
	row.partition = p
	return row
}

// GetRow retrieves a specific row from this Partition
func (p *partitionImpl) GetRow(rowNum int) sif.Row {
	return &rowImpl{rowNum, p}
}

func (p *partitionImpl) GetSchema() sif.Schema {
	return p.schema
}

// FromStreamedData loads data from a protobuf stream into this Partition
func FromStreamedData(stream pb.PartitionsService_TransferPartitionDataClient, partitionMeta *pb.MPartitionMeta, schema sif.Schema, compressor sif.Compressor) (sif.Partition, error) {
	// stream data for Partition
	dataOffset := 0
	buff := make([]byte, partitionMeta.GetBytes())
	for chunk, err := stream.Recv(); err != io.EOF; chunk, err = stream.Recv() {
		if err != nil {
			// not an EOF, but something else
			return nil, err
		}
		switch chunk.DataType {
		case iutil.SerializedPartitionDataType:
			copy(buff[dataOffset:dataOffset+len(chunk.Data)], chunk.Data)
			dataOffset += len(chunk.Data)
		default:
			return nil, fmt.Errorf("unknown chunk data type encountered: %d", chunk.DataType)
		}
	}
	// confirm we received the correct amount of data
	if uint32(dataOffset) != partitionMeta.GetBytes() {
		return nil, fmt.Errorf("streamed %d bytes for SerializedPartition %s. Expected %d", dataOffset, partitionMeta.GetId(), partitionMeta.GetBytes())
	}
	// decompress and deserialize
	return FromBytes(buff, schema, compressor)
}