mirror of https://github.com/grafana/loki
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
462 lines
12 KiB
462 lines
12 KiB
package array_test
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"math"
|
|
"math/rand"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/grafana/loki/v3/pkg/columnar"
|
|
"github.com/grafana/loki/v3/pkg/columnar/columnartest"
|
|
"github.com/grafana/loki/v3/pkg/columnar/types"
|
|
"github.com/grafana/loki/v3/pkg/dataset/array"
|
|
"github.com/grafana/loki/v3/pkg/dataset/buffer"
|
|
"github.com/grafana/loki/v3/pkg/memory"
|
|
)
|
|
|
|
func TestZigZagCodec_Validation(t *testing.T) {
|
|
data := &array.SpecPlain{}
|
|
|
|
tt := []struct {
|
|
name string
|
|
spec array.Spec
|
|
typ types.Type
|
|
|
|
expectError bool
|
|
}{
|
|
{
|
|
name: "accepts non-nullable int32",
|
|
spec: &array.SpecZigZag{Data: data},
|
|
typ: &types.Int32{Nullable: false},
|
|
expectError: false,
|
|
},
|
|
{
|
|
name: "accepts non-nullable int64",
|
|
spec: &array.SpecZigZag{Data: data},
|
|
typ: &types.Int64{Nullable: false},
|
|
expectError: false,
|
|
},
|
|
{
|
|
name: "accepts nullable int64 with validity in data spec",
|
|
spec: &array.SpecZigZag{Data: &array.SpecPlain{Validity: &array.SpecBool{}}},
|
|
typ: &types.Int64{Nullable: true},
|
|
expectError: false,
|
|
},
|
|
{
|
|
name: "rejects uint32",
|
|
spec: &array.SpecZigZag{Data: data},
|
|
typ: &types.Uint32{Nullable: false},
|
|
expectError: true,
|
|
},
|
|
{
|
|
name: "rejects uint64",
|
|
spec: &array.SpecZigZag{Data: data},
|
|
typ: &types.Uint64{Nullable: false},
|
|
expectError: true,
|
|
},
|
|
{
|
|
name: "rejects bool",
|
|
spec: &array.SpecZigZag{Data: data},
|
|
typ: &types.Bool{Nullable: false},
|
|
expectError: true,
|
|
},
|
|
{
|
|
name: "rejects utf8",
|
|
spec: &array.SpecZigZag{Data: data},
|
|
typ: &types.UTF8{Nullable: false},
|
|
expectError: true,
|
|
},
|
|
{
|
|
name: "rejects nil data spec",
|
|
spec: &array.SpecZigZag{Data: nil},
|
|
typ: &types.Int64{Nullable: false},
|
|
expectError: true,
|
|
},
|
|
}
|
|
|
|
for _, tc := range tt {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
var alloc memory.Allocator
|
|
_, err := array.NewWriter(&alloc, tc.spec, tc.typ)
|
|
|
|
if tc.expectError {
|
|
require.Error(t, err)
|
|
} else {
|
|
require.NoError(t, err)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestZigZagCodec_NonNullable(t *testing.T) {
|
|
t.Run("int32", func(t *testing.T) {
|
|
var alloc memory.Allocator
|
|
var store buffer.MemoryStore
|
|
|
|
var (
|
|
spec = &array.SpecZigZag{Data: &array.SpecPlain{}}
|
|
typ = &types.Int32{Nullable: false}
|
|
)
|
|
|
|
w, err := array.NewWriter(&alloc, spec, typ)
|
|
require.NoError(t, err)
|
|
|
|
input := []columnar.Array{
|
|
columnartest.Array(t, types.KindInt32, &alloc, int32(-3), int32(-2), int32(-1), int32(0), int32(1)),
|
|
columnartest.Array(t, types.KindInt32, &alloc, int32(2), int32(3), int32(100), int32(-100)),
|
|
}
|
|
for _, arr := range input {
|
|
require.NoError(t, w.Append(arr))
|
|
}
|
|
|
|
result, err := w.Flush(t.Context(), &store)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 9, result.RowCount)
|
|
require.Equal(t, 0, result.Stats.NullCount)
|
|
|
|
r, err := array.NewReader(&alloc, result, &store)
|
|
require.NoError(t, err)
|
|
|
|
expect := columnartest.Array(
|
|
t, types.KindInt32, &alloc,
|
|
int32(-3), int32(-2), int32(-1), int32(0), int32(1),
|
|
int32(2), int32(3), int32(100), int32(-100),
|
|
)
|
|
|
|
actual := readBatches(t, &alloc, r, 2) // Read in a small batch size to test reading multiple times
|
|
columnartest.RequireArraysEqual(t, expect, actual, memory.Bitmap{})
|
|
|
|
_, err = r.Read(t.Context(), &alloc, 1)
|
|
require.ErrorIs(t, err, io.EOF)
|
|
})
|
|
|
|
t.Run("int64", func(t *testing.T) {
|
|
var alloc memory.Allocator
|
|
var store buffer.MemoryStore
|
|
|
|
var (
|
|
spec = &array.SpecZigZag{Data: &array.SpecPlain{}}
|
|
typ = &types.Int64{Nullable: false}
|
|
)
|
|
|
|
w, err := array.NewWriter(&alloc, spec, typ)
|
|
require.NoError(t, err)
|
|
|
|
input := []columnar.Array{
|
|
columnartest.Array(t, types.KindInt64, &alloc, int64(-3), int64(-2), int64(-1), int64(0), int64(1)),
|
|
columnartest.Array(t, types.KindInt64, &alloc, int64(2), int64(3), int64(100), int64(-100)),
|
|
}
|
|
for _, arr := range input {
|
|
require.NoError(t, w.Append(arr))
|
|
}
|
|
|
|
result, err := w.Flush(t.Context(), &store)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 9, result.RowCount)
|
|
require.Equal(t, 0, result.Stats.NullCount)
|
|
|
|
r, err := array.NewReader(&alloc, result, &store)
|
|
require.NoError(t, err)
|
|
|
|
expect := columnartest.Array(
|
|
t, types.KindInt64, &alloc,
|
|
int64(-3), int64(-2), int64(-1), int64(0), int64(1),
|
|
int64(2), int64(3), int64(100), int64(-100),
|
|
)
|
|
|
|
actual := readBatches(t, &alloc, r, 2) // Read in a small batch size to test reading multiple times
|
|
columnartest.RequireArraysEqual(t, expect, actual, memory.Bitmap{})
|
|
|
|
_, err = r.Read(t.Context(), &alloc, 1)
|
|
require.ErrorIs(t, err, io.EOF)
|
|
})
|
|
}
|
|
|
|
func TestZigZagCodec_Nullable(t *testing.T) {
|
|
var alloc memory.Allocator
|
|
var store buffer.MemoryStore
|
|
|
|
var (
|
|
spec = &array.SpecZigZag{Data: &array.SpecPlain{Validity: &array.SpecBool{}}}
|
|
typ = &types.Int64{Nullable: true}
|
|
)
|
|
|
|
w, err := array.NewWriter(&alloc, spec, typ)
|
|
require.NoError(t, err)
|
|
|
|
input := []columnar.Array{
|
|
columnartest.Array(t, types.KindInt64, &alloc, int64(-1), nil, int64(1), nil, int64(0)),
|
|
columnartest.Array(t, types.KindInt64, &alloc, int64(42), int64(-42)),
|
|
columnartest.Array(t, types.KindInt64, &alloc, nil),
|
|
}
|
|
for _, arr := range input {
|
|
require.NoError(t, w.Append(arr))
|
|
}
|
|
|
|
result, err := w.Flush(t.Context(), &store)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 8, result.RowCount)
|
|
require.Equal(t, 3, result.Stats.NullCount)
|
|
|
|
r, err := array.NewReader(&alloc, result, &store)
|
|
require.NoError(t, err)
|
|
|
|
expect := columnartest.Array(
|
|
t, types.KindInt64, &alloc,
|
|
int64(-1), nil, int64(1), nil, int64(0),
|
|
int64(42), int64(-42),
|
|
nil,
|
|
)
|
|
|
|
actual := readBatches(t, &alloc, r, 2) // Read in a small batch size to test reading multiple times
|
|
columnartest.RequireArraysEqual(t, expect, actual, memory.Bitmap{})
|
|
|
|
_, err = r.Read(t.Context(), &alloc, 1)
|
|
require.ErrorIs(t, err, io.EOF)
|
|
}
|
|
|
|
func TestZigZagCodec_PartialRead(t *testing.T) {
|
|
var alloc memory.Allocator
|
|
var store buffer.MemoryStore
|
|
|
|
var (
|
|
spec = &array.SpecZigZag{Data: &array.SpecPlain{}}
|
|
typ = &types.Int64{Nullable: false}
|
|
)
|
|
|
|
w, err := array.NewWriter(&alloc, spec, typ)
|
|
require.NoError(t, err)
|
|
|
|
input := columnartest.Array(t, types.KindInt64, &alloc,
|
|
int64(-5), int64(-4), int64(-3), int64(-2), int64(-1),
|
|
int64(0), int64(1), int64(2), int64(3), int64(4),
|
|
)
|
|
require.NoError(t, w.Append(input))
|
|
|
|
result, err := w.Flush(t.Context(), &store)
|
|
require.NoError(t, err)
|
|
|
|
r, err := array.NewReader(&alloc, result, &store)
|
|
require.NoError(t, err)
|
|
|
|
batch1, err := r.Read(t.Context(), &alloc, 3)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 3, batch1.Len())
|
|
columnartest.RequireArraysEqual(t,
|
|
columnartest.Array(t, types.KindInt64, &alloc, int64(-5), int64(-4), int64(-3)),
|
|
batch1, memory.Bitmap{},
|
|
)
|
|
|
|
batch2, err := r.Read(t.Context(), &alloc, 3)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 3, batch2.Len())
|
|
columnartest.RequireArraysEqual(t,
|
|
columnartest.Array(t, types.KindInt64, &alloc, int64(-2), int64(-1), int64(0)),
|
|
batch2, memory.Bitmap{},
|
|
)
|
|
|
|
batch3, err := r.Read(t.Context(), &alloc, 3)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 3, batch3.Len())
|
|
columnartest.RequireArraysEqual(t,
|
|
columnartest.Array(t, types.KindInt64, &alloc, int64(1), int64(2), int64(3)),
|
|
batch3, memory.Bitmap{},
|
|
)
|
|
|
|
// Last batch: only 1 remaining.
|
|
batch4, err := r.Read(t.Context(), &alloc, 3)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 1, batch4.Len())
|
|
columnartest.RequireArraysEqual(t,
|
|
columnartest.Array(t, types.KindInt64, &alloc, int64(4)),
|
|
batch4, memory.Bitmap{},
|
|
)
|
|
|
|
_, err = r.Read(t.Context(), &alloc, 1)
|
|
require.ErrorIs(t, err, io.EOF)
|
|
}
|
|
|
|
func TestZigZagCodec_Extremes(t *testing.T) {
|
|
t.Run("int32", func(t *testing.T) {
|
|
var alloc memory.Allocator
|
|
var store buffer.MemoryStore
|
|
|
|
var (
|
|
spec = &array.SpecZigZag{Data: &array.SpecPlain{}}
|
|
typ = &types.Int32{Nullable: false}
|
|
)
|
|
|
|
w, err := array.NewWriter(&alloc, spec, typ)
|
|
require.NoError(t, err)
|
|
|
|
input := columnartest.Array(t, types.KindInt32, &alloc,
|
|
int32(0), int32(math.MinInt32), int32(math.MaxInt32), int32(-1), int32(1),
|
|
)
|
|
require.NoError(t, w.Append(input))
|
|
|
|
result, err := w.Flush(t.Context(), &store)
|
|
require.NoError(t, err)
|
|
|
|
r, err := array.NewReader(&alloc, result, &store)
|
|
require.NoError(t, err)
|
|
|
|
actual := readBatches(t, &alloc, r, 2) // Read in a small batch size to test reading multiple times
|
|
columnartest.RequireArraysEqual(t, input, actual, memory.Bitmap{})
|
|
})
|
|
|
|
t.Run("int64", func(t *testing.T) {
|
|
var alloc memory.Allocator
|
|
var store buffer.MemoryStore
|
|
|
|
var (
|
|
spec = &array.SpecZigZag{Data: &array.SpecPlain{}}
|
|
typ = &types.Int64{Nullable: false}
|
|
)
|
|
|
|
w, err := array.NewWriter(&alloc, spec, typ)
|
|
require.NoError(t, err)
|
|
|
|
input := columnartest.Array(t, types.KindInt64, &alloc,
|
|
int64(0), int64(math.MinInt64), int64(math.MaxInt64), int64(-1), int64(1),
|
|
)
|
|
require.NoError(t, w.Append(input))
|
|
|
|
result, err := w.Flush(t.Context(), &store)
|
|
require.NoError(t, err)
|
|
|
|
r, err := array.NewReader(&alloc, result, &store)
|
|
require.NoError(t, err)
|
|
|
|
actual := readBatches(t, &alloc, r, 2) // Read in a small batch size to test reading multiple times
|
|
columnartest.RequireArraysEqual(t, input, actual, memory.Bitmap{})
|
|
})
|
|
}
|
|
|
|
func TestZigZagCodec_WithBitpacked(t *testing.T) {
|
|
var alloc memory.Allocator
|
|
var store buffer.MemoryStore
|
|
|
|
var (
|
|
spec = &array.SpecZigZag{
|
|
Data: &array.SpecBitpacked{BlockSize: 8, Widths: &array.SpecPlain{}},
|
|
}
|
|
typ = &types.Int64{Nullable: false}
|
|
)
|
|
|
|
w, err := array.NewWriter(&alloc, spec, typ)
|
|
require.NoError(t, err)
|
|
|
|
input := columnartest.Array(t, types.KindInt64, &alloc,
|
|
int64(0), int64(-1), int64(1), int64(-2), int64(2),
|
|
int64(-3), int64(3), int64(-4), int64(4), int64(-5),
|
|
)
|
|
require.NoError(t, w.Append(input))
|
|
|
|
result, err := w.Flush(t.Context(), &store)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 10, result.RowCount)
|
|
|
|
r, err := array.NewReader(&alloc, result, &store)
|
|
require.NoError(t, err)
|
|
|
|
actual := readBatches(t, &alloc, r, 2) // Read in a small batch size to test reading multiple times
|
|
columnartest.RequireArraysEqual(t, input, actual, memory.Bitmap{})
|
|
}
|
|
|
|
func BenchmarkZigZagCodec(b *testing.B) {
|
|
var (
|
|
store buffer.MemoryStore
|
|
|
|
spec = &array.SpecZigZag{Data: &array.SpecPlain{}}
|
|
typ = &types.Int64{Nullable: false}
|
|
)
|
|
|
|
const valuesPerPage = 1 << 16
|
|
|
|
type scenario struct {
|
|
name string
|
|
valueCount int
|
|
encoded array.Array
|
|
}
|
|
|
|
build := func(name string, valueCount int, valueAt func(i int) int64) scenario {
|
|
var alloc memory.Allocator
|
|
w, err := array.NewWriter(&alloc, spec, typ)
|
|
require.NoError(b, err)
|
|
|
|
builder := columnar.NewNumberBuilder[int64](&alloc)
|
|
builder.Grow(valueCount)
|
|
|
|
for i := range valueCount {
|
|
builder.AppendValue(valueAt(i))
|
|
}
|
|
require.NoError(b, w.Append(builder.Build()))
|
|
|
|
arr, err := w.Flush(b.Context(), &store)
|
|
require.NoError(b, err)
|
|
return scenario{name: name, valueCount: valueCount, encoded: arr}
|
|
}
|
|
|
|
scenarios := []scenario{
|
|
build("variance=constant", valuesPerPage, func(int) int64 { return -42 }),
|
|
build("variance=small_range", valuesPerPage, func(i int) int64 { return int64(i%16) - 8 }),
|
|
func() scenario {
|
|
rnd := rand.New(rand.NewSource(0))
|
|
return build("variance=random", valuesPerPage, func(int) int64 {
|
|
return rnd.Int63() - rnd.Int63()
|
|
})
|
|
}(),
|
|
}
|
|
|
|
batchSizes := []int{256, 1024, 4096}
|
|
|
|
for _, sc := range scenarios {
|
|
b.Run(sc.name, func(b *testing.B) {
|
|
b.Run(fmt.Sprintf("values_per_page=%d", sc.valueCount), func(b *testing.B) {
|
|
for _, batchSize := range batchSizes {
|
|
b.Run(fmt.Sprintf("batch_size=%d", batchSize), func(b *testing.B) {
|
|
var alloc memory.Allocator
|
|
|
|
b.ReportAllocs()
|
|
decodedBytesPerOp := int64(sc.valueCount) * 8
|
|
b.SetBytes(decodedBytesPerOp)
|
|
|
|
for b.Loop() {
|
|
alloc.Reset()
|
|
|
|
r, _ := array.NewReader(&alloc, sc.encoded, &store)
|
|
|
|
var decoded int
|
|
for {
|
|
arr, err := r.Read(b.Context(), &alloc, batchSize)
|
|
if arr != nil {
|
|
decoded += arr.Len()
|
|
}
|
|
|
|
if errors.Is(err, io.EOF) {
|
|
break
|
|
} else if err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
}
|
|
|
|
if decoded != sc.valueCount {
|
|
b.Fatalf("decoded %d values, expected %d", decoded, sc.valueCount)
|
|
}
|
|
}
|
|
|
|
elapsed := b.Elapsed()
|
|
if elapsed > 0 {
|
|
totalDecoded := int64(sc.valueCount) * int64(b.N)
|
|
b.ReportMetric(float64(totalDecoded)/elapsed.Seconds(), "rows/s")
|
|
}
|
|
})
|
|
}
|
|
})
|
|
})
|
|
}
|
|
}
|
|
|