mirror of
https://github.com/hanzoai/traces.git
synced 2026-10-30 18:20:41 +00:00
cc86e98b04
* Add operators IN, NOT IN, MATCH ANY, and MATCH NONE
* Implement Static.append()
* Print string arrays with '`' instead of '"'
Example ["aa", "bb"] -> [`aa`, `bb`]
* Rewrite binary operations to array equivalent
The rewriter handles the following cases:
- { .a = "a" || .a = "b" } => { .a IN ["a", "b"] }
- { .a != "a" && .a != "b" } => { .a NOT IN ["a", "b"] }
- { .a =~ "a" || .a =~ "b" } => { .a MATCH ANY ["a", "b"] }
- { .a !~ "a" && .a !~ "b" } => { .a MATCH NONE ["a", "b"] }
* Add test cases with arrays to TestBinOp
* traceql.Parse() applies AST optimizations by default
* Rename predicate_test.go to predicates_test.go
* Add array predicates to parquetquery
* BREAKING CHANGE: more precise regex validation
BinaryOperation.validate() considered everythin a valid regex where
FieldExpression.EncodeToString(false) happened to return a valid regex.
With this change the RHS of a regex operation is required to be a string
or string array
* vp4: handle predicates with arrays correctly
* vp5: handle predicates with arrays correctly
* vp3: handle predicates with arrays correctly
* Add test cases for optimized queries
* Update query test examples with optimizable queries
* vp3: fix handling of array types
* CHANGELOG.md
* Pass array operators down to fetch layer
* Read to first virtual row after seeking on a low definition level
This helps avoid additional spans/op in queries with only
resource attributes
---------
Co-authored-by: Martin Disibio <[email protected]>
471 lines
12 KiB
Go
471 lines
12 KiB
Go
package parquetquery
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/parquet-go/parquet-go"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
type testInt struct {
|
|
N int64 `parquet:",dict"`
|
|
}
|
|
|
|
type mockPredicate struct {
|
|
ret bool
|
|
valCalled bool
|
|
pageCalled bool
|
|
chunkCalled bool
|
|
}
|
|
|
|
type testDictString struct {
|
|
S string `parquet:",dict"`
|
|
}
|
|
|
|
var _ Predicate = (*mockPredicate)(nil)
|
|
|
|
func newAlwaysTruePredicate() *mockPredicate {
|
|
return &mockPredicate{ret: true}
|
|
}
|
|
|
|
func newAlwaysFalsePredicate() *mockPredicate {
|
|
return &mockPredicate{ret: false}
|
|
}
|
|
|
|
func (p *mockPredicate) String() string { return "mockPredicate{}" }
|
|
func (p *mockPredicate) KeepValue(parquet.Value) bool { p.valCalled = true; return p.ret }
|
|
func (p *mockPredicate) KeepPage(parquet.Page) bool { p.pageCalled = true; return p.ret }
|
|
func (p *mockPredicate) KeepColumnChunk(*ColumnChunkHelper) bool { p.chunkCalled = true; return p.ret }
|
|
|
|
type predicateTestCase struct {
|
|
testName string
|
|
writeData func(w *parquet.Writer) //nolint:all
|
|
keptChunks int
|
|
keptPages int
|
|
keptValues int
|
|
predicate Predicate
|
|
}
|
|
|
|
func TestSubstringPredicate(t *testing.T) {
|
|
testCases := []predicateTestCase{
|
|
{
|
|
testName: "all chunks/pages/values inspected",
|
|
predicate: NewSubstringPredicate("b"),
|
|
keptChunks: 1,
|
|
keptPages: 1,
|
|
keptValues: 2,
|
|
writeData: func(w *parquet.Writer) { //nolint:all
|
|
|
|
require.NoError(t, w.Write(&testDictString{"abc"})) // kept
|
|
require.NoError(t, w.Write(&testDictString{"bcd"})) // kept
|
|
require.NoError(t, w.Write(&testDictString{"cde"})) // skipped
|
|
},
|
|
},
|
|
{
|
|
testName: "dictionary in the page header allows for skipping a page",
|
|
predicate: NewSubstringPredicate("x"), // Not present in any values
|
|
keptChunks: 0,
|
|
keptPages: 0,
|
|
keptValues: 0,
|
|
writeData: func(w *parquet.Writer) { //nolint:all
|
|
require.NoError(t, w.Write(&testDictString{"abc"}))
|
|
require.NoError(t, w.Write(&testDictString{"abc"}))
|
|
require.NoError(t, w.Write(&testDictString{"abc"}))
|
|
require.NoError(t, w.Write(&testDictString{"abc"}))
|
|
require.NoError(t, w.Write(&testDictString{"abc"}))
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tC := range testCases {
|
|
t.Run(tC.testName, func(t *testing.T) {
|
|
testPredicate(t, tC)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestNewStringInPredicate(t *testing.T) {
|
|
testCases := []predicateTestCase{
|
|
{
|
|
testName: "all chunks/pages/values inspected",
|
|
predicate: func() Predicate {
|
|
return NewStringInPredicate([]string{"abc", "acd"})
|
|
}(),
|
|
keptChunks: 1,
|
|
keptPages: 1,
|
|
keptValues: 2,
|
|
writeData: func(w *parquet.Writer) { //nolint:all
|
|
require.NoError(t, w.Write(&testDictString{"abc"})) // kept
|
|
require.NoError(t, w.Write(&testDictString{"acd"})) // kept
|
|
require.NoError(t, w.Write(&testDictString{"cde"})) // skipped
|
|
},
|
|
},
|
|
{
|
|
testName: "dictionary in the page header allows for skipping a column chunk",
|
|
predicate: func() Predicate {
|
|
return NewStringInPredicate([]string{"x"})
|
|
}(), // Not present in any values
|
|
keptChunks: 0,
|
|
keptPages: 0,
|
|
keptValues: 0,
|
|
writeData: func(w *parquet.Writer) { //nolint:all
|
|
require.NoError(t, w.Write(&testDictString{"abc"}))
|
|
require.NoError(t, w.Write(&testDictString{"abc"}))
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tC := range testCases {
|
|
t.Run(tC.testName, func(t *testing.T) {
|
|
testPredicate(t, tC)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestNewRegexInPredicate(t *testing.T) {
|
|
testCases := []predicateTestCase{
|
|
{
|
|
testName: "all chunks/pages/values inspected",
|
|
predicate: func() Predicate {
|
|
pred, err := NewRegexInPredicate([]string{"a.*"})
|
|
require.NoError(t, err)
|
|
|
|
return pred
|
|
}(),
|
|
keptChunks: 1,
|
|
keptPages: 1,
|
|
keptValues: 2,
|
|
writeData: func(w *parquet.Writer) { //nolint:all
|
|
require.NoError(t, w.Write(&testDictString{"abc"})) // kept
|
|
require.NoError(t, w.Write(&testDictString{"acd"})) // kept
|
|
require.NoError(t, w.Write(&testDictString{"cde"})) // skipped
|
|
},
|
|
},
|
|
{
|
|
testName: "dictionary in the page header allows for skipping a column chunk",
|
|
predicate: func() Predicate {
|
|
pred, err := NewRegexInPredicate([]string{"x.*"})
|
|
require.NoError(t, err)
|
|
|
|
return pred
|
|
}(), // Not present in any values
|
|
keptChunks: 0,
|
|
keptPages: 0,
|
|
keptValues: 0,
|
|
writeData: func(w *parquet.Writer) { //nolint:all
|
|
require.NoError(t, w.Write(&testDictString{"abc"}))
|
|
require.NoError(t, w.Write(&testDictString{"abc"}))
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tC := range testCases {
|
|
t.Run(tC.testName, func(t *testing.T) {
|
|
testPredicate(t, tC)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestNewStringNotInPredicate(t *testing.T) {
|
|
testCases := []predicateTestCase{
|
|
{
|
|
testName: "all chunks/pages/values inspected",
|
|
predicate: func() Predicate {
|
|
return NewStringNotInPredicate([]string{"abc", "acd"})
|
|
}(),
|
|
keptChunks: 1,
|
|
keptPages: 1,
|
|
keptValues: 2,
|
|
writeData: func(w *parquet.Writer) { //nolint:all
|
|
require.NoError(t, w.Write(&testDictString{"abc"})) // skipped
|
|
require.NoError(t, w.Write(&testDictString{"acd"})) // skipped
|
|
require.NoError(t, w.Write(&testDictString{"cde"})) // kept
|
|
require.NoError(t, w.Write(&testDictString{"xde"})) // kept
|
|
},
|
|
},
|
|
{
|
|
testName: "dictionary in the page header allows for skipping a column chunk",
|
|
predicate: func() Predicate {
|
|
return NewStringNotInPredicate([]string{"xyz"})
|
|
}(), // All values excluded
|
|
keptChunks: 0,
|
|
keptPages: 0,
|
|
keptValues: 0,
|
|
writeData: func(w *parquet.Writer) { //nolint:all
|
|
require.NoError(t, w.Write(&testDictString{"xyz"}))
|
|
require.NoError(t, w.Write(&testDictString{"xyz"}))
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tC := range testCases {
|
|
t.Run(tC.testName, func(t *testing.T) {
|
|
testPredicate(t, tC)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestNewRegexNotInPredicate(t *testing.T) {
|
|
testCases := []predicateTestCase{
|
|
{
|
|
testName: "all chunks/pages/values inspected",
|
|
predicate: func() Predicate {
|
|
pred, err := NewRegexNotInPredicate([]string{"a.*"})
|
|
require.NoError(t, err)
|
|
|
|
return pred
|
|
}(),
|
|
keptChunks: 1,
|
|
keptPages: 1,
|
|
keptValues: 2,
|
|
writeData: func(w *parquet.Writer) { //nolint:all
|
|
require.NoError(t, w.Write(&testDictString{"abc"})) // skipped
|
|
require.NoError(t, w.Write(&testDictString{"acd"})) // skipped
|
|
require.NoError(t, w.Write(&testDictString{"cde"})) // kept
|
|
require.NoError(t, w.Write(&testDictString{"xde"})) // kept
|
|
},
|
|
},
|
|
{
|
|
testName: "dictionary in the page header allows for skipping a column chunk",
|
|
predicate: func() Predicate {
|
|
pred, err := NewRegexNotInPredicate([]string{"x.*"})
|
|
require.NoError(t, err)
|
|
|
|
return pred
|
|
}(), // Not present in any values
|
|
keptChunks: 0,
|
|
keptPages: 0,
|
|
keptValues: 0,
|
|
writeData: func(w *parquet.Writer) { //nolint:all
|
|
require.NoError(t, w.Write(&testDictString{"xyz"}))
|
|
require.NoError(t, w.Write(&testDictString{"xyz"}))
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tC := range testCases {
|
|
t.Run(tC.testName, func(t *testing.T) {
|
|
testPredicate(t, tC)
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestOrPredicateCallsKeepColumnChunk ensures that the OrPredicate calls
|
|
// KeepColumnChunk on all of its children. This is important because the
|
|
// Dictionary predicates rely on KeepColumnChunk always being called at the
|
|
// beginning of a row group to reset their page.
|
|
func TestOrPredicateCallsKeepColumnChunk(t *testing.T) {
|
|
tcs := []struct {
|
|
preds []*mockPredicate
|
|
}{
|
|
{},
|
|
{
|
|
preds: []*mockPredicate{
|
|
newAlwaysTruePredicate(),
|
|
},
|
|
},
|
|
{
|
|
preds: []*mockPredicate{
|
|
newAlwaysFalsePredicate(),
|
|
},
|
|
},
|
|
{
|
|
preds: []*mockPredicate{
|
|
newAlwaysFalsePredicate(),
|
|
newAlwaysTruePredicate(),
|
|
},
|
|
},
|
|
{
|
|
preds: []*mockPredicate{
|
|
newAlwaysTruePredicate(),
|
|
newAlwaysFalsePredicate(),
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tc := range tcs {
|
|
preds := make([]Predicate, 0, len(tc.preds)+1)
|
|
for _, pred := range tc.preds {
|
|
preds = append(preds, pred)
|
|
}
|
|
|
|
recordPred := &mockPredicate{}
|
|
preds = append(preds, recordPred)
|
|
|
|
p := NewOrPredicate(preds...)
|
|
p.KeepColumnChunk(nil)
|
|
|
|
for _, pred := range preds {
|
|
require.True(t, pred.(*mockPredicate).chunkCalled)
|
|
}
|
|
}
|
|
}
|
|
|
|
// testPredicate by writing data and then iterating the column.
|
|
// The data model must contain a single column.
|
|
func testPredicate(t *testing.T, tc predicateTestCase) {
|
|
t.Helper()
|
|
buf := new(bytes.Buffer)
|
|
w := parquet.NewWriter(buf)
|
|
tc.writeData(w)
|
|
w.Flush()
|
|
w.Close()
|
|
|
|
file := bytes.NewReader(buf.Bytes())
|
|
r, err := parquet.OpenFile(file, int64(buf.Len()))
|
|
require.NoError(t, err)
|
|
|
|
p := InstrumentedPredicate{Pred: tc.predicate}
|
|
|
|
i := NewSyncIterator(context.TODO(), r.RowGroups(), 0, SyncIteratorOptPredicate(&p))
|
|
for {
|
|
res, err := i.Next()
|
|
require.NoError(t, err)
|
|
if res == nil {
|
|
break
|
|
}
|
|
}
|
|
|
|
require.Equal(t, tc.keptChunks, int(p.KeptColumnChunks), "keptChunks")
|
|
require.Equal(t, tc.keptPages, int(p.KeptPages), "keptPages")
|
|
require.Equal(t, tc.keptValues, int(p.KeptValues), "keptValues")
|
|
}
|
|
|
|
func BenchmarkSubstringPredicate(b *testing.B) {
|
|
p := NewSubstringPredicate("abc")
|
|
|
|
s := make([]parquet.Value, 1000)
|
|
for i := 0; i < 1000; i++ {
|
|
s[i] = parquet.ValueOf(uuid.New().String())
|
|
}
|
|
|
|
b.ResetTimer()
|
|
|
|
for i := 0; i < b.N; i++ {
|
|
for _, ss := range s {
|
|
p.KeepValue(ss)
|
|
}
|
|
}
|
|
}
|
|
|
|
func BenchmarkStringInPredicate(b *testing.B) {
|
|
p := NewStringInPredicate([]string{"abc"})
|
|
|
|
s := make([]parquet.Value, 1000)
|
|
for i := 0; i < 1000; i++ {
|
|
s[i] = parquet.ValueOf(uuid.New().String())
|
|
}
|
|
|
|
b.ResetTimer()
|
|
|
|
for i := 0; i < b.N; i++ {
|
|
for _, ss := range s {
|
|
p.KeepValue(ss)
|
|
}
|
|
}
|
|
}
|
|
|
|
func BenchmarkRegexInPredicate(b *testing.B) {
|
|
p, err := NewRegexInPredicate([]string{"abc"})
|
|
require.NoError(b, err)
|
|
|
|
s := make([]parquet.Value, 1000)
|
|
for i := 0; i < 1000; i++ {
|
|
s[i] = parquet.ValueOf(uuid.New().String())
|
|
}
|
|
|
|
b.ResetTimer()
|
|
|
|
for i := 0; i < b.N; i++ {
|
|
for _, ss := range s {
|
|
p.KeepValue(ss)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestIntInPredicate(t *testing.T) {
|
|
testCases := []predicateTestCase{
|
|
{
|
|
testName: "all chunks/pages/values inspected",
|
|
predicate: func() Predicate {
|
|
var p Predicate = NewIntInPredicate([]int64{1, 3})
|
|
return p
|
|
}(),
|
|
keptChunks: 1,
|
|
keptPages: 1,
|
|
keptValues: 2,
|
|
writeData: func(w *parquet.Writer) { //nolint:all
|
|
require.NoError(t, w.Write(&testInt{1})) // kept
|
|
require.NoError(t, w.Write(&testInt{2})) // skipped
|
|
require.NoError(t, w.Write(&testInt{3})) // kept
|
|
require.NoError(t, w.Write(&testInt{10})) // skipped
|
|
},
|
|
},
|
|
{
|
|
testName: "dictionary allows skipping a column chunk when no ints match",
|
|
predicate: func() Predicate {
|
|
var p Predicate = NewIntInPredicate([]int64{0, 4, 100})
|
|
return p
|
|
}(),
|
|
keptChunks: 0,
|
|
keptPages: 0,
|
|
keptValues: 0,
|
|
writeData: func(w *parquet.Writer) { //nolint:all
|
|
require.NoError(t, w.Write(&testInt{1}))
|
|
require.NoError(t, w.Write(&testInt{2}))
|
|
require.NoError(t, w.Write(&testInt{3}))
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tC := range testCases {
|
|
t.Run(tC.testName, func(t *testing.T) {
|
|
testPredicate(t, tC)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestIntNotInPredicate(t *testing.T) {
|
|
testCases := []predicateTestCase{
|
|
{
|
|
testName: "all chunks/pages/values inspected",
|
|
predicate: func() Predicate {
|
|
var p Predicate = NewIntNotInPredicate([]int64{1, 3})
|
|
return p
|
|
}(),
|
|
keptChunks: 1,
|
|
keptPages: 1,
|
|
keptValues: 2,
|
|
writeData: func(w *parquet.Writer) { //nolint:all
|
|
require.NoError(t, w.Write(&testInt{1})) // skipped
|
|
require.NoError(t, w.Write(&testInt{2})) // kept
|
|
require.NoError(t, w.Write(&testInt{3})) // skipped
|
|
require.NoError(t, w.Write(&testInt{4})) // kept
|
|
},
|
|
},
|
|
{
|
|
testName: "dictionary allows skipping a column chunk when all ints excluded",
|
|
predicate: func() Predicate {
|
|
var p Predicate = NewIntNotInPredicate([]int64{7, 8})
|
|
return p
|
|
}(),
|
|
keptChunks: 0,
|
|
keptPages: 0,
|
|
keptValues: 0,
|
|
writeData: func(w *parquet.Writer) { //nolint:all
|
|
require.NoError(t, w.Write(&testInt{7}))
|
|
require.NoError(t, w.Write(&testInt{8}))
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tC := range testCases {
|
|
t.Run(tC.testName, func(t *testing.T) {
|
|
testPredicate(t, tC)
|
|
})
|
|
}
|
|
}
|