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.
300 lines
8.5 KiB
300 lines
8.5 KiB
package layout_test
|
|
|
|
import (
|
|
"io"
|
|
"math"
|
|
"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/compute"
|
|
"github.com/grafana/loki/v3/pkg/dataset/array"
|
|
"github.com/grafana/loki/v3/pkg/dataset/buffer"
|
|
"github.com/grafana/loki/v3/pkg/dataset/layout"
|
|
"github.com/grafana/loki/v3/pkg/expr"
|
|
"github.com/grafana/loki/v3/pkg/memory"
|
|
)
|
|
|
|
func TestCodec_Array(t *testing.T) {
|
|
var alloc memory.Allocator
|
|
var store buffer.MemoryStore
|
|
|
|
spec := &layout.SpecArray{
|
|
Spec: &array.SpecBool{Validity: nil},
|
|
}
|
|
|
|
w, err := layout.NewWriter(&alloc, &store, spec, &types.Bool{Nullable: false})
|
|
require.NoError(t, err)
|
|
|
|
expect := columnartest.Array(t, types.KindBool, &alloc, true, false, true, true)
|
|
require.NoError(t, w.Append(t.Context(), expect))
|
|
|
|
l, err := w.Flush(t.Context())
|
|
require.NoError(t, err)
|
|
|
|
r, err := layout.NewReader(&alloc, l, &store)
|
|
require.NoError(t, err)
|
|
|
|
// Advance the reader to get everything.
|
|
n, err := r.Next(t.Context(), math.MaxInt)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 4, n)
|
|
actual, err := r.Project(t.Context(), &alloc, &expr.Identity{}, memory.Bitmap{})
|
|
require.NoError(t, err)
|
|
|
|
columnartest.RequireArraysEqual(t, expect, actual, memory.Bitmap{})
|
|
}
|
|
|
|
func TestCodec_Array_Windowing(t *testing.T) {
|
|
var alloc memory.Allocator
|
|
var store buffer.MemoryStore
|
|
|
|
spec := &layout.SpecArray{
|
|
Spec: &array.SpecPlain{Validity: nil},
|
|
}
|
|
|
|
w, err := layout.NewWriter(&alloc, &store, spec, &types.Int64{Nullable: false})
|
|
require.NoError(t, err)
|
|
|
|
input := columnartest.Array(t, types.KindInt64, &alloc,
|
|
int64(10), int64(20), int64(30), int64(40), int64(50), int64(60),
|
|
)
|
|
require.NoError(t, w.Append(t.Context(), input))
|
|
|
|
l, err := w.Flush(t.Context())
|
|
require.NoError(t, err)
|
|
|
|
r, err := layout.NewReader(&alloc, l, &store)
|
|
require.NoError(t, err)
|
|
|
|
// First window: 2 rows.
|
|
n, err := r.Next(t.Context(), 2)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 2, n)
|
|
|
|
actual, err := r.Project(t.Context(), &alloc, &expr.Identity{}, memory.Bitmap{})
|
|
require.NoError(t, err)
|
|
columnartest.RequireArraysEqual(t,
|
|
columnartest.Array(t, types.KindInt64, &alloc, int64(10), int64(20)),
|
|
actual,
|
|
memory.Bitmap{},
|
|
)
|
|
|
|
// Second window: 2 rows.
|
|
n, err = r.Next(t.Context(), 2)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 2, n)
|
|
|
|
actual, err = r.Project(t.Context(), &alloc, &expr.Identity{}, memory.Bitmap{})
|
|
require.NoError(t, err)
|
|
columnartest.RequireArraysEqual(t,
|
|
columnartest.Array(t, types.KindInt64, &alloc, int64(30), int64(40)),
|
|
actual,
|
|
memory.Bitmap{},
|
|
)
|
|
|
|
// Third window: 2 remaining rows.
|
|
n, err = r.Next(t.Context(), 2)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 2, n)
|
|
|
|
actual, err = r.Project(t.Context(), &alloc, &expr.Identity{}, memory.Bitmap{})
|
|
require.NoError(t, err)
|
|
columnartest.RequireArraysEqual(t,
|
|
columnartest.Array(t, types.KindInt64, &alloc, int64(50), int64(60)),
|
|
actual,
|
|
memory.Bitmap{},
|
|
)
|
|
|
|
// Fourth call: EOF.
|
|
n, err = r.Next(t.Context(), 2)
|
|
require.ErrorIs(t, err, io.EOF)
|
|
require.Equal(t, 0, n)
|
|
}
|
|
|
|
func TestCodec_Array_Filter(t *testing.T) {
|
|
var alloc memory.Allocator
|
|
var store buffer.MemoryStore
|
|
|
|
spec := &layout.SpecArray{
|
|
Spec: &array.SpecPlain{Validity: nil},
|
|
}
|
|
|
|
w, err := layout.NewWriter(&alloc, &store, spec, &types.Int64{Nullable: false})
|
|
require.NoError(t, err)
|
|
|
|
input := columnartest.Array(t, types.KindInt64, &alloc,
|
|
int64(10), int64(20), int64(30), int64(40),
|
|
)
|
|
require.NoError(t, w.Append(t.Context(), input))
|
|
|
|
l, err := w.Flush(t.Context())
|
|
require.NoError(t, err)
|
|
|
|
r, err := layout.NewReader(&alloc, l, &store)
|
|
require.NoError(t, err)
|
|
|
|
n, err := r.Next(t.Context(), math.MaxInt)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 4, n)
|
|
|
|
// Filter: Identity > 20 → keeps rows where value > 20.
|
|
filterExpr := &expr.Binary{
|
|
Left: &expr.Identity{},
|
|
Op: expr.BinaryOpGT,
|
|
Right: &expr.Constant{Value: &columnar.NumberScalar[int64]{Value: 20}},
|
|
}
|
|
|
|
mask, err := r.Filter(t.Context(), &alloc, filterExpr, memory.Bitmap{})
|
|
require.NoError(t, err)
|
|
|
|
// Project with the filter mask; Project returns full window, so
|
|
// apply compute.Filter to compact.
|
|
projected, err := r.Project(t.Context(), &alloc, &expr.Identity{}, mask)
|
|
require.NoError(t, err)
|
|
filtered, err := compute.Filter(&alloc, projected, mask)
|
|
require.NoError(t, err)
|
|
columnartest.RequireArraysEqual(t,
|
|
columnartest.Array(t, types.KindInt64, &alloc, int64(30), int64(40)),
|
|
filtered.(columnar.Array),
|
|
memory.Bitmap{},
|
|
)
|
|
}
|
|
|
|
func TestCodec_Array_Filter_WithInputMask(t *testing.T) {
|
|
var alloc memory.Allocator
|
|
var store buffer.MemoryStore
|
|
|
|
spec := &layout.SpecArray{
|
|
Spec: &array.SpecPlain{Validity: nil},
|
|
}
|
|
|
|
w, err := layout.NewWriter(&alloc, &store, spec, &types.Int64{Nullable: false})
|
|
require.NoError(t, err)
|
|
|
|
input := columnartest.Array(t, types.KindInt64, &alloc,
|
|
int64(10), int64(20), int64(30), int64(40),
|
|
)
|
|
require.NoError(t, w.Append(t.Context(), input))
|
|
|
|
l, err := w.Flush(t.Context())
|
|
require.NoError(t, err)
|
|
|
|
r, err := layout.NewReader(&alloc, l, &store)
|
|
require.NoError(t, err)
|
|
|
|
n, err := r.Next(t.Context(), math.MaxInt)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 4, n)
|
|
|
|
// Filter: Identity > 15 → would match rows 1,2,3 (values 20,30,40).
|
|
// But provide an input mask that only includes rows 0 and 2.
|
|
filterExpr := &expr.Binary{
|
|
Left: &expr.Identity{},
|
|
Op: expr.BinaryOpGT,
|
|
Right: &expr.Constant{Value: &columnar.NumberScalar[int64]{Value: 15}},
|
|
}
|
|
|
|
inputMask := memory.NewBitmap(&alloc, 4)
|
|
inputMask.AppendValues(true, false, true, false)
|
|
|
|
mask, err := r.Filter(t.Context(), &alloc, filterExpr, inputMask)
|
|
require.NoError(t, err)
|
|
|
|
// Project with the intersected mask → only row 2 (value 30).
|
|
projected, err := r.Project(t.Context(), &alloc, &expr.Identity{}, mask)
|
|
require.NoError(t, err)
|
|
filtered, err := compute.Filter(&alloc, projected, mask)
|
|
require.NoError(t, err)
|
|
columnartest.RequireArraysEqual(t,
|
|
columnartest.Array(t, types.KindInt64, &alloc, int64(30)),
|
|
filtered.(columnar.Array),
|
|
memory.Bitmap{},
|
|
)
|
|
}
|
|
|
|
func TestCodec_Array_Filter_NullsDoNotLeak(t *testing.T) {
|
|
var alloc memory.Allocator
|
|
var store buffer.MemoryStore
|
|
|
|
spec := &layout.SpecArray{
|
|
Spec: &array.SpecPlain{Validity: &array.SpecBool{}},
|
|
}
|
|
|
|
w, err := layout.NewWriter(&alloc, &store, spec, &types.Int64{Nullable: true})
|
|
require.NoError(t, err)
|
|
|
|
// Rows 1 and 3 are null. NumberBuilder.AppendNull writes the zero value
|
|
// for the backing storage, so the predicate "Identity > -5" evaluates to
|
|
// true at those positions in the compute kernel's values bitmap.
|
|
input := columnartest.Array(t, types.KindInt64, &alloc,
|
|
int64(10), nil, int64(30), nil,
|
|
)
|
|
require.NoError(t, w.Append(t.Context(), input))
|
|
|
|
l, err := w.Flush(t.Context())
|
|
require.NoError(t, err)
|
|
|
|
r, err := layout.NewReader(&alloc, l, &store)
|
|
require.NoError(t, err)
|
|
|
|
n, err := r.Next(t.Context(), math.MaxInt)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 4, n)
|
|
|
|
// Identity > -5: backing zero for nulls also satisfies this.
|
|
filterExpr := &expr.Binary{
|
|
Left: &expr.Identity{},
|
|
Op: expr.BinaryOpGT,
|
|
Right: &expr.Constant{Value: &columnar.NumberScalar[int64]{Value: -5}},
|
|
}
|
|
|
|
mask, err := r.Filter(t.Context(), &alloc, filterExpr, memory.Bitmap{})
|
|
require.NoError(t, err)
|
|
|
|
require.Equal(t, 4, mask.Len())
|
|
require.True(t, mask.Get(0), "row 0 (value 10) should pass")
|
|
require.False(t, mask.Get(1), "row 1 (null) must not pass")
|
|
require.True(t, mask.Get(2), "row 2 (value 30) should pass")
|
|
require.False(t, mask.Get(3), "row 3 (null) must not pass")
|
|
}
|
|
|
|
func TestCodec_Array_ProjectAllRows(t *testing.T) {
|
|
var alloc memory.Allocator
|
|
var store buffer.MemoryStore
|
|
|
|
spec := &layout.SpecArray{
|
|
Spec: &array.SpecPlain{Validity: nil},
|
|
}
|
|
|
|
w, err := layout.NewWriter(&alloc, &store, spec, &types.Int64{Nullable: false})
|
|
require.NoError(t, err)
|
|
|
|
input := columnartest.Array(t, types.KindInt64, &alloc,
|
|
int64(1), int64(2), int64(3),
|
|
)
|
|
require.NoError(t, w.Append(t.Context(), input))
|
|
|
|
l, err := w.Flush(t.Context())
|
|
require.NoError(t, err)
|
|
|
|
r, err := layout.NewReader(&alloc, l, &store)
|
|
require.NoError(t, err)
|
|
|
|
n, err := r.Next(t.Context(), math.MaxInt)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 3, n)
|
|
|
|
// Project with an all-set mask should return everything.
|
|
allSet := memory.NewBitmap(&alloc, 3)
|
|
allSet.AppendCount(true, 3)
|
|
|
|
projected, err := r.Project(t.Context(), &alloc, &expr.Identity{}, allSet)
|
|
require.NoError(t, err)
|
|
filtered, err := compute.Filter(&alloc, projected, allSet)
|
|
require.NoError(t, err)
|
|
columnartest.RequireArraysEqual(t, input, filtered.(columnar.Array), memory.Bitmap{})
|
|
}
|
|
|