perf: Optimization pass for reading from dataobjs (#16747)

Author: benclive

pkg/dataobj/builder.go

@@ -195,8 +195,8 @@ func (b *Builder) Append(stream logproto.Stream) error {
 		b.logs.Append(logs.Record{
 			StreamID:  streamID,
 			Timestamp: entry.Timestamp,
-			Metadata:  entry.StructuredMetadata,
-			Line:      entry.Line,
+			Metadata:  convertMetadata(entry.StructuredMetadata),
+			Line:      []byte(entry.Line),
 		})
 	}
 

pkg/dataobj/internal/dataset/page_builder.go

@@ -175,6 +175,9 @@ func valueSize(v Value) int {
 		// Assuming that strings are PLAIN encoded using their length and bytes.
 		str := v.String()
 		return binary.Size(len(str)) + len(str)
+	case datasetmd.VALUE_TYPE_BYTE_ARRAY:
+		arr := v.ByteArray()
+		return binary.Size(len(arr)) + len(arr)
 	}
 
 	return 0

pkg/dataobj/internal/dataset/value.go

@@ -1,6 +1,7 @@
 package dataset
 
 import (
+	"bytes"
 	"cmp"
 	"encoding/binary"
 	"fmt"
@@ -12,6 +13,7 @@ import (
 // Helper types
 type (
 	stringptr *byte
+	bytearray *byte
 )
 
 // A Value represents a single value within a dataset. Unlike [any], Values can
@@ -69,6 +71,14 @@ func StringValue(v string) Value {
 	}
 }
 
+// ByteArrayValue returns a [Value] for a byte slice representing a string.
+func ByteArrayValue(v []byte) Value {
+	return Value{
+		num: uint64(len(v)),
+		any: (bytearray)(unsafe.SliceData(v)),
+	}
+}
+
 // IsNil returns whether v is nil.
 func (v Value) IsNil() bool {
 	return v.any == nil
@@ -93,6 +103,8 @@ func (v Value) Type() datasetmd.ValueType {
 		return v
 	case stringptr:
 		return datasetmd.VALUE_TYPE_STRING
+	case bytearray:
+		return datasetmd.VALUE_TYPE_BYTE_ARRAY
 	default:
 		panic(fmt.Sprintf("dataset.Value has unexpected type %T", v))
 	}
@@ -126,6 +138,16 @@ func (v Value) String() string {
 	return v.Type().String()
 }
 
+// ByteSlice returns v's value as a byte slice. If v is not a string,
+// ByteSlice returns a byte slice of the form "VALUE_TYPE_T", where T is the
+// underlying type of v.
+func (v Value) ByteArray() []byte {
+	if ba, ok := v.any.(bytearray); ok {
+		return unsafe.Slice(ba, v.num)
+	}
+	panic(fmt.Sprintf("dataset.Value type is %s, not %s", v.Type(), datasetmd.VALUE_TYPE_BYTE_ARRAY))
+}
+
 // MarshalBinary encodes v into a binary representation. Non-NULL values encode
 // first with the type (encoded as uvarint), followed by an encoded value,
 // where:
@@ -150,6 +172,8 @@ func (v Value) MarshalBinary() (data []byte, err error) {
 	case datasetmd.VALUE_TYPE_STRING:
 		str := v.String()
 		buf = append(buf, unsafe.Slice(unsafe.StringData(str), len(str))...)
+	case datasetmd.VALUE_TYPE_BYTE_ARRAY:
+		buf = append(buf, v.ByteArray()...)
 	default:
 		return nil, fmt.Errorf("dataset.Value.MarshalBinary: unsupported type %s", v.Type())
 	}
@@ -186,6 +210,8 @@ func (v *Value) UnmarshalBinary(data []byte) error {
 	case datasetmd.VALUE_TYPE_STRING:
 		str := string(data[n:])
 		*v = StringValue(str)
+	case datasetmd.VALUE_TYPE_BYTE_ARRAY:
+		*v = ByteArrayValue(data[n:])
 	default:
 		return fmt.Errorf("dataset.Value.UnmarshalBinary: unsupported type %s", vtyp)
 	}
@@ -221,6 +247,8 @@ func CompareValues(a, b Value) int {
 		return cmp.Compare(a.Uint64(), b.Uint64())
 	case datasetmd.VALUE_TYPE_STRING:
 		return cmp.Compare(a.String(), b.String())
+	case datasetmd.VALUE_TYPE_BYTE_ARRAY:
+		return bytes.Compare(a.ByteArray(), b.ByteArray())
 	default:
 		panic(fmt.Sprintf("page.CompareValues: unsupported type %s", a.Type()))
 	}

pkg/dataobj/internal/dataset/value_encoding_plain.go

@@ -16,35 +16,41 @@ func init() {
 	registerValueEncoding(
 		datasetmd.VALUE_TYPE_STRING,
 		datasetmd.ENCODING_TYPE_PLAIN,
-		func(w streamio.Writer) valueEncoder { return newPlainEncoder(w) },
-		func(r streamio.Reader) valueDecoder { return newPlainDecoder(r) },
+		func(w streamio.Writer) valueEncoder { return newPlainStringEncoder(w) },
+		func(r streamio.Reader) valueDecoder { return newPlainStringDecoder(r) },
+	)
+	registerValueEncoding(
+		datasetmd.VALUE_TYPE_BYTE_ARRAY,
+		datasetmd.ENCODING_TYPE_PLAIN,
+		func(w streamio.Writer) valueEncoder { return newPlainBytesEncoder(w) },
+		func(r streamio.Reader) valueDecoder { return newPlainBytesDecoder(r) },
 	)
 }
 
-// A plainEncoder encodes string values to an [streamio.Writer].
-type plainEncoder struct {
+// A plainStringEncoder encodes string values to an [streamio.Writer].
+type plainStringEncoder struct {
 	w streamio.Writer
 }
 
-var _ valueEncoder = (*plainEncoder)(nil)
+var _ valueEncoder = (*plainStringEncoder)(nil)
 
 // newPlainEncoder creates a plainEncoder that writes encoded strings to w.
-func newPlainEncoder(w streamio.Writer) *plainEncoder {
-	return &plainEncoder{w: w}
+func newPlainStringEncoder(w streamio.Writer) *plainStringEncoder {
+	return &plainStringEncoder{w: w}
 }
 
 // ValueType returns [datasetmd.VALUE_TYPE_STRING].
-func (enc *plainEncoder) ValueType() datasetmd.ValueType {
+func (enc *plainStringEncoder) ValueType() datasetmd.ValueType {
 	return datasetmd.VALUE_TYPE_STRING
 }
 
 // EncodingType returns [datasetmd.ENCODING_TYPE_PLAIN].
-func (enc *plainEncoder) EncodingType() datasetmd.EncodingType {
+func (enc *plainStringEncoder) EncodingType() datasetmd.EncodingType {
 	return datasetmd.ENCODING_TYPE_PLAIN
 }
 
 // Encode encodes an individual string value.
-func (enc *plainEncoder) Encode(v Value) error {
+func (enc *plainStringEncoder) Encode(v Value) error {
 	if v.Type() != datasetmd.VALUE_TYPE_STRING {
 		return fmt.Errorf("plain: invalid value type %v", v.Type())
 	}
@@ -65,41 +71,41 @@ func (enc *plainEncoder) Encode(v Value) error {
 }
 
 // Flush implements [valueEncoder]. It is a no-op for plainEncoder.
-func (enc *plainEncoder) Flush() error {
+func (enc *plainStringEncoder) Flush() error {
 	return nil
 }
 
 // Reset implements [valueEncoder]. It resets the encoder to write to w.
-func (enc *plainEncoder) Reset(w streamio.Writer) {
+func (enc *plainStringEncoder) Reset(w streamio.Writer) {
 	enc.w = w
 }
 
-// plainDecoder decodes strings from an [streamio.Reader].
-type plainDecoder struct {
+// plainStringDecoder decodes strings from an [streamio.Reader].
+type plainStringDecoder struct {
 	r streamio.Reader
 }
 
-var _ valueDecoder = (*plainDecoder)(nil)
+var _ valueDecoder = (*plainStringDecoder)(nil)
 
-// newPlainDecoder creates a plainDecoder that reads encoded strings from r.
-func newPlainDecoder(r streamio.Reader) *plainDecoder {
-	return &plainDecoder{r: r}
+// newPlainStringDecoder creates a plainDecoder that reads encoded strings from r.
+func newPlainStringDecoder(r streamio.Reader) *plainStringDecoder {
+	return &plainStringDecoder{r: r}
 }
 
-// ValueType returns [datasetmd.VALUE_TYPE_STRING].
-func (dec *plainDecoder) ValueType() datasetmd.ValueType {
+// ValueType returns [datasetmd.VALUE_TYPE_BYTE_ARRAY].
+func (dec *plainStringDecoder) ValueType() datasetmd.ValueType {
 	return datasetmd.VALUE_TYPE_STRING
 }
 
 // EncodingType returns [datasetmd.ENCODING_TYPE_PLAIN].
-func (dec *plainDecoder) EncodingType() datasetmd.EncodingType {
+func (dec *plainStringDecoder) EncodingType() datasetmd.EncodingType {
 	return datasetmd.ENCODING_TYPE_PLAIN
 }
 
 // Decode decodes up to len(s) values, storing the results into s. The
 // number of decoded values is returned, followed by an error (if any).
 // At the end of the stream, Decode returns 0, [io.EOF].
-func (dec *plainDecoder) Decode(s []Value) (int, error) {
+func (dec *plainStringDecoder) Decode(s []Value) (int, error) {
 	if len(s) == 0 {
 		return 0, nil
 	}
@@ -123,7 +129,7 @@ func (dec *plainDecoder) Decode(s []Value) (int, error) {
 }
 
 // decode decodes a string.
-func (dec *plainDecoder) decode() (Value, error) {
+func (dec *plainStringDecoder) decode() (Value, error) {
 	sz, err := binary.ReadUvarint(dec.r)
 	if err != nil {
 		return StringValue(""), err
@@ -137,6 +143,123 @@ func (dec *plainDecoder) decode() (Value, error) {
 }
 
 // Reset implements [valueDecoder]. It resets the decoder to read from r.
-func (dec *plainDecoder) Reset(r streamio.Reader) {
+func (dec *plainStringDecoder) Reset(r streamio.Reader) {
+	dec.r = r
+}
+
+// A plainBytesEncoder encodes byte array values to an [streamio.Writer].
+type plainBytesEncoder struct {
+	w streamio.Writer
+}
+
+var _ valueEncoder = (*plainStringEncoder)(nil)
+
+// newPlainEncoder creates a plainEncoder that writes encoded strings to w.
+func newPlainBytesEncoder(w streamio.Writer) *plainBytesEncoder {
+	return &plainBytesEncoder{w: w}
+}
+
+// ValueType returns [datasetmd.VALUE_TYPE_BYTE_ARRAY].
+func (enc *plainBytesEncoder) ValueType() datasetmd.ValueType {
+	return datasetmd.VALUE_TYPE_BYTE_ARRAY
+}
+
+// EncodingType returns [datasetmd.ENCODING_TYPE_PLAIN].
+func (enc *plainBytesEncoder) EncodingType() datasetmd.EncodingType {
+	return datasetmd.ENCODING_TYPE_PLAIN
+}
+
+// Encode encodes an individual string value.
+func (enc *plainBytesEncoder) Encode(v Value) error {
+	if v.Type() != datasetmd.VALUE_TYPE_BYTE_ARRAY {
+		return fmt.Errorf("plain: invalid value type %v", v.Type())
+	}
+	sv := v.ByteArray()
+
+	if err := streamio.WriteUvarint(enc.w, uint64(len(sv))); err != nil {
+		return err
+	}
+
+	n, err := enc.w.Write(sv)
+	if n != len(sv) {
+		return fmt.Errorf("short write; expected %d bytes, wrote %d", len(sv), n)
+	}
+	return err
+}
+
+// Flush implements [valueEncoder]. It is a no-op for plainEncoder.
+func (enc *plainBytesEncoder) Flush() error {
+	return nil
+}
+
+// Reset implements [valueEncoder]. It resets the encoder to write to w.
+func (enc *plainBytesEncoder) Reset(w streamio.Writer) {
+	enc.w = w
+}
+
+// plainBytesDecoder decodes byte arrays from an [streamio.Reader].
+type plainBytesDecoder struct {
+	r streamio.Reader
+}
+
+var _ valueDecoder = (*plainBytesDecoder)(nil)
+
+// newPlainBytesDecoder creates a plainDecoder that reads encoded strings from r.
+func newPlainBytesDecoder(r streamio.Reader) *plainBytesDecoder {
+	return &plainBytesDecoder{r: r}
+}
+
+// ValueType returns [datasetmd.VALUE_TYPE_BYTE_ARRAY].
+func (dec *plainBytesDecoder) ValueType() datasetmd.ValueType {
+	return datasetmd.VALUE_TYPE_BYTE_ARRAY
+}
+
+// EncodingType returns [datasetmd.ENCODING_TYPE_PLAIN].
+func (dec *plainBytesDecoder) EncodingType() datasetmd.EncodingType {
+	return datasetmd.ENCODING_TYPE_PLAIN
+}
+
+// Decode decodes up to len(s) values, storing the results into s. The
+// number of decoded values is returned, followed by an error (if any).
+// At the end of the stream, Decode returns 0, [io.EOF].
+func (dec *plainBytesDecoder) Decode(s []Value) (int, error) {
+	if len(s) == 0 {
+		return 0, nil
+	}
+
+	var err error
+	var v Value
+
+	for i := range s {
+		v, err = dec.decode()
+		if errors.Is(err, io.EOF) {
+			if i == 0 {
+				return 0, io.EOF
+			}
+			return i, nil
+		} else if err != nil {
+			return i, err
+		}
+		s[i] = v
+	}
+	return len(s), nil
+}
+
+// decode decodes a string.
+func (dec *plainBytesDecoder) decode() (Value, error) {
+	sz, err := binary.ReadUvarint(dec.r)
+	if err != nil {
+		return ByteArrayValue([]byte{}), err
+	}
+
+	dst := make([]byte, int(sz))
+	if _, err := io.ReadFull(dec.r, dst); err != nil {
+		return ByteArrayValue([]byte{}), err
+	}
+	return ByteArrayValue(dst), nil
+}
+
+// Reset implements [valueDecoder]. It resets the decoder to read from r.
+func (dec *plainBytesDecoder) Reset(r streamio.Reader) {
 	dec.r = r
 }