Repository: polarsignals/frostdb
Pull request #737 (2024-02-13 10:02:56)

Fix logicalplan validations


Patch:
diff --git a/query/logicalplan/logicalplan.go b/query/logicalplan/logicalplan.go
index a0fc4dc..25a72bc 100644
--- a/query/logicalplan/logicalplan.go
+++ b/query/logicalplan/logicalplan.go
@@ -125,6 +125,12 @@ func (plan *LogicalPlan) DataTypeForExpr(expr Expr) (arrow.DataType, error) {
 
 		return t, nil
 	case plan.Projection != nil:
+		for _, e := range plan.Projection.Exprs {
+			if e.Name() == expr.Name() {
+				return e.DataType(plan.Input)
+			}
+		}
+
 		t, err := expr.DataType(plan.Input)
 		if err != nil {
 			return nil, fmt.Errorf("data type for expr %v within Projection: %w", expr, err)
@@ -255,10 +261,13 @@ func (scan *TableScan) DataTypeForExpr(expr Expr) (arrow.DataType, error) {
 	}
 
 	s := tp.Schema()
+	if s == nil {
+		return nil, fmt.Errorf("table %q has no schema", scan.TableName)
+	}
 
 	t, err := DataTypeForExprWithSchema(expr, s)
 	if err != nil {
-		return nil, fmt.Errorf("type for expr %q in table %q", expr, scan.TableName)
+		return nil, fmt.Errorf("type for expr %q in table %q: %w", expr, scan.TableName, err)
 	}
 
 	return t, nil
@@ -373,7 +382,7 @@ type Projection struct {
 }
 
 func (p *Projection) String() string {
-	return "Projection"
+	return "Projection (" + fmt.Sprint(p.Exprs) + ")"
 }
 
 type Aggregation struct {
diff --git a/query/logicalplan/validate.go b/query/logicalplan/validate.go
index 11193db..2b80ee4 100644
--- a/query/logicalplan/validate.go
+++ b/query/logicalplan/validate.go
@@ -5,6 +5,7 @@ import (
 	"reflect"
 	"strings"
 
+	"github.com/apache/arrow/go/v14/arrow"
 	"github.com/apache/arrow/go/v14/arrow/scalar"
 	"github.com/parquet-go/parquet-go/format"
 )
@@ -67,9 +68,9 @@ func Validate(plan *LogicalPlan) error {
 	if err == nil {
 		switch {
 		case plan.SchemaScan != nil:
-			err = nil
+			err = ValidateSchemaScan(plan)
 		case plan.TableScan != nil:
-			err = nil
+			err = ValidateTableScan(plan)
 		case plan.Filter != nil:
 			err = ValidateFilter(plan)
 		case plan.Distinct != nil:
@@ -143,6 +144,86 @@ func ValidateSingleFieldSet(plan *LogicalPlan) *PlanValidationError {
 	return nil
 }
 
+func ValidateSchemaScan(plan *LogicalPlan) *PlanValidationError {
+	if plan.SchemaScan.TableProvider == nil {
+		return &PlanValidationError{
+			plan:    plan,
+			message: "table provider must not be nil",
+		}
+	}
+
+	if plan.SchemaScan.TableName == "" {
+		return &PlanValidationError{
+			plan:    plan,
+			message: "table name must not be empty",
+		}
+	}
+
+	tableReader, err := plan.SchemaScan.TableProvider.GetTable(plan.SchemaScan.TableName)
+	if err != nil {
+		return &PlanValidationError{
+			plan:    plan,
+			message: fmt.Sprintf("failed to get table: %s", err),
+		}
+	}
+	if tableReader == nil {
+		return &PlanValidationError{
+			plan:    plan,
+			message: "table not found",
+		}
+	}
+
+	schema := tableReader.Schema()
+	if schema == nil {
+		return &PlanValidationError{
+			plan:    plan,
+			message: "table schema must not be nil",
+		}
+	}
+
+	return nil
+}
+
+func ValidateTableScan(plan *LogicalPlan) *PlanValidationError {
+	if plan.TableScan.TableProvider == nil {
+		return &PlanValidationError{
+			plan:    plan,
+			message: "table provider must not be nil",
+		}
+	}
+
+	if plan.TableScan.TableName == "" {
+		return &PlanValidationError{
+			plan:    plan,
+			message: "table name must not be empty",
+		}
+	}
+
+	tableReader, err := plan.TableScan.TableProvider.GetTable(plan.TableScan.TableName)
+	if err != nil {
+		return &PlanValidationError{
+			plan:    plan,
+			message: fmt.Sprintf("failed to get table: %s", err),
+		}
+	}
+	if tableReader == nil {
+		return &PlanValidationError{
+			plan:    plan,
+			message: "table not found",
+		}
+	}
+
+	schema := tableReader.Schema()
+	if schema == nil {
+		return &PlanValidationError{
+			plan:    plan,
+			message: "table schema must not be nil",
+		}
+	}
+
+	return nil
+}
+
 // ValidateAggregation validates the logical plan's aggregation step.
 func ValidateAggregation(plan *LogicalPlan) *PlanValidationError {
 	// check that the expression is not nil
@@ -172,60 +253,28 @@ type Named interface {
 
 func ValidateAggregationExpr(plan *LogicalPlan) *ExprValidationError {
 	for _, expr := range plan.Aggregation.AggExprs {
-		// check that the aggregation expression has the required structure
-		colFinder := newTypeFinder((*Column)(nil))
-		expr.Accept(&colFinder)
-
-		dynColFinder := newTypeFinder((*DynamicColumn)(nil))
-		expr.Accept(&dynColFinder)
-
-		aggFuncFinder := newTypeFinder((*AggregationFunction)(nil))
-		expr.Accept(&aggFuncFinder)
-
-		if (colFinder.result == nil && dynColFinder.result == nil) || aggFuncFinder.result == nil {
+		t, err := expr.Expr.DataType(plan.Input)
+		if err != nil {
 			return &ExprValidationError{
-				message: "aggregation expression is invalid. must contain AggregationFunction and Column",
-				expr:    expr,
+				expr:    expr.Expr,
+				message: fmt.Errorf("get type of expression to aggregate: %w", err).Error(),
 			}
 		}
 
-		// check that column being aggregated on exists in the schema
-		schema := plan.InputSchema()
-		if schema == nil {
-			return nil // cannot check column type if there's no input schema
-		}
-
-		var named Named
-		named = colFinder.result
-		if named == nil {
-			named = dynColFinder.result
-		}
-
-		name, _, _ := strings.Cut(named.Name(), ".")
-		column, found := schema.ColumnByName(name)
-		if !found {
+		if t == nil {
 			return &ExprValidationError{
-				message: fmt.Sprintf("column not found: %s", named.Name()),
-				expr:    expr,
+				expr:    expr.Expr,
+				message: "invalid aggregation: expression type cannot be determined",
 			}
 		}
 
-		// check that the column type can be aggregated by the function type
-		columnType := column.StorageLayout.Type()
-		aggFuncExpr := aggFuncFinder.result.(*AggregationFunction)
-		logicalType := columnType.LogicalType()
-		if logicalType != nil && logicalType.UTF8 != nil {
-			switch aggFuncExpr.Func {
-			case AggFuncSum:
-				return &ExprValidationError{
-					message: "cannot sum text column",
-					expr:    expr,
-				}
-			case AggFuncMax:
-				return &ExprValidationError{
-					message: "cannot max text column",
-					expr:    expr,
-				}
+		switch t {
+		case arrow.PrimitiveTypes.Int64, arrow.PrimitiveTypes.Float64:
+			// valid
+		default:
+			return &ExprValidationError{
+				expr:    expr.Expr,
+				message: fmt.Errorf("invalid aggregation: expression type %s is not supported", t).Error(),
 			}
 		}
 	}


Test patch:
diff --git a/db_test.go b/db_test.go
index d13d21a..009cd38 100644
--- a/db_test.go
+++ b/db_test.go
@@ -24,7 +24,6 @@ import (
 	"google.golang.org/protobuf/proto"
 
 	"github.com/polarsignals/frostdb/dynparquet"
-	schemapb "github.com/polarsignals/frostdb/gen/proto/go/frostdb/schema/v1alpha1"
 	walpb "github.com/polarsignals/frostdb/gen/proto/go/frostdb/wal/v1alpha1"
 	"github.com/polarsignals/frostdb/index"
 	"github.com/polarsignals/frostdb/query"
@@ -356,6 +355,10 @@ func Test_DB_WithStorage(t *testing.T) {
 
 	db, err = c.DB(context.Background(), t.Name())
 	require.NoError(t, err)
+
+	_, err = db.Table(t.Name(), config)
+	require.NoError(t, err)
+
 	engine = query.NewEngine(pool, db.TableProvider())
 	var onDisk arrow.Record
 	err = engine.ScanTable(t.Name()).
@@ -372,253 +375,6 @@ func Test_DB_WithStorage(t *testing.T) {
 	require.Equal(t, int64(2), onDisk.NumRows())
 }
 
-func Test_DB_ColdStart(t *testing.T) {
-	sanitize := func(name string) string {
-		return strings.Replace(name, "/", "-", -1)
-	}
-
-	config := NewTableConfig(
-		dynparquet.SampleDefinition(),
-	)
-
-	bucket := objstore.NewInMemBucket()
-	sinksource := NewDefaultObjstoreBucket(bucket)
-	logger := newTestLogger(t)
-
-	tests := map[string]struct {
-		newColumnstore func(t *testing.T) *ColumnStore
-	}{
-		"cold start with storage": {
-			newColumnstore: func(t *testing.T) *ColumnStore {
-				c, err := New(
-					WithLogger(logger),
-					WithReadWriteStorage(sinksource),
-				)
-				require.NoError(t, err)
-				return c
-			},
-		},
-		"cold start with storage and wal": {
-			newColumnstore: func(t *testing.T) *ColumnStore {
-				c, err := New(
-					WithLogger(logger),
-					WithReadWriteStorage(sinksource),
-					WithWAL(),
-					WithStoragePath(t.TempDir()),
-				)
-				require.NoError(t, err)
-				return c
-			},
-		},
-	}
-
-	for name, test := range tests {
-		t.Run(name, func(t *testing.T) {
-			c := test.newColumnstore(t)
-			db, err := c.DB(context.Background(), sanitize(t.Name()))
-			require.NoError(t, err)
-			table, err := db.Table(sanitize(t.Name()), config)
-			require.NoError(t, err)
-			t.Cleanup(func() {
-				os.RemoveAll(sanitize(t.Name()))
-			})
-
-			samples := dynparquet.Samples{
-				{
-					ExampleType: "test",
-					Labels: map[string]string{
-						"label1": "value1",
-						"label2": "value2",
-					},
-					Stacktrace: []uuid.UUID{
-						{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1},
-						{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x2},
-					},
-					Timestamp: 1,
-					Value:     1,
-				},
-				{
-					ExampleType: "test",
-					Labels: map[string]string{
-						"label1": "value1",
-						"label2": "value2",
-					},
-					Stacktrace: []uuid.UUID{
-						{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1},
-						{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x2},
-					},
-					Timestamp: 2,
-					Value:     2,
-				},
-				{
-					ExampleType: "test",
-					Labels: map[string]string{
-						"label1": "value1",
-						"label2": "value2",
-					},
-					Stacktrace: []uuid.UUID{
-						{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1},
-						{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x2},
-					},
-					Timestamp: 3,
-					Value:     3,
-				},
-			}
-
-			r, err := samples.ToRecord()
-			require.NoError(t, err)
-
-			ctx := context.Background()
-			_, err = table.InsertRecord(ctx, r)
-			require.NoError(t, err)
-
-			// Gracefully close the db to persist blocks
-			c.Close()
-
-			// Open a new database pointed to the same bucket storage
-			c, err = New(
-				WithLogger(logger),
-				WithReadWriteStorage(sinksource),
-			)
-			require.NoError(t, err)
-			defer c.Close()
-
-			// connect to our test db
-			db, err = c.DB(context.Background(), sanitize(t.Name()))
-			require.NoError(t, err)
-
-			pool := memory.NewGoAllocator()
-			engine := query.NewEngine(pool, db.TableProvider())
-			require.NoError(t, engine.ScanTable(sanitize(t.Name())).Execute(
-				context.Background(), func(ctx context.Context, r arrow.Record) error {
-					require.Equal(t, int64(6), r.NumCols())
-					require.Equal(t, int64(3), r.NumRows())
-					return nil
-				},
-			))
-		})
-	}
-}
-
-func Test_DB_ColdStart_MissingColumn(t *testing.T) {
-	schemaDef := &schemapb.Schema{
-		Name: "test",
-		Columns: []*schemapb.Column{
-			{
-				Name: "example_type",
-				StorageLayout: &schemapb.StorageLayout{
-					Type:     schemapb.StorageLayout_TYPE_STRING,
-					Encoding: schemapb.StorageLayout_ENCODING_RLE_DICTIONARY,
-				},
-				Dynamic: false,
-			},
-			{
-				Name: "labels",
-				StorageLayout: &schemapb.StorageLayout{
-					Type:     schemapb.StorageLayout_TYPE_STRING,
-					Nullable: true,
-					Encoding: schemapb.StorageLayout_ENCODING_RLE_DICTIONARY,
-				},
-				Dynamic: true,
-			},
-			{
-				Name: "pprof_labels",
-				StorageLayout: &schemapb.StorageLayout{
-					Type:     schemapb.StorageLayout_TYPE_STRING,
-					Nullable: true,
-					Encoding: schemapb.StorageLayout_ENCODING_RLE_DICTIONARY,
-				},
-				Dynamic: true,
-			},
-		},
-		SortingColumns: []*schemapb.SortingColumn{
-			{
-				Name:      "example_type",
-				Direction: schemapb.SortingColumn_DIRECTION_ASCENDING,
-			},
-			{
-				Name:       "labels",
-				Direction:  schemapb.SortingColumn_DIRECTION_ASCENDING,
-				NullsFirst: true,
-			},
-			{
-				Name:       "pprof_labels",
-				Direction:  schemapb.SortingColumn_DIRECTION_ASCENDING,
-				NullsFirst: true,
-			},
-		},
-	}
-
-	config := NewTableConfig(schemaDef)
-
-	bucket := objstore.NewInMemBucket()
-
-	sinksource := NewDefaultObjstoreBucket(bucket)
-	logger := newTestLogger(t)
-
-	c, err := New(
-		WithLogger(logger),
-		WithReadWriteStorage(sinksource),
-	)
-	require.NoError(t, err)
-
-	db, err := c.DB(context.Background(), t.Name())
-	require.NoError(t, err)
-	table, err := db.Table(t.Name(), config)
-	require.NoError(t, err)
-	t.Cleanup(func() {
-		os.RemoveAll(t.Name())
-	})
-
-	schema := arrow.NewSchema([]arrow.Field{
-		{Name: "example_type", Type: arrow.BinaryTypes.Binary},
-		{Name: "labels.label1", Type: arrow.BinaryTypes.Binary},
-		{Name: "labels.label2", Type: arrow.BinaryTypes.Binary},
-	}, nil)
-	bldr := array.NewRecordBuilder(memory.DefaultAllocator, schema)
-	defer bldr.Release()
-
-	bldr.Field(0).(*array.BinaryBuilder).Append([]byte("test"))
-	bldr.Field(1).(*array.BinaryBuilder).Append([]byte("value1"))
-	bldr.Field(2).(*array.BinaryBuilder).Append([]byte("value1"))
-
-	r := bldr.NewRecord()
-	defer r.Release()
-
-	ctx := context.Background()
-	_, err = table.InsertRecord(ctx, r)
-	require.NoError(t, err)
-
-	// Gracefully close the db to persist blocks
-	c.Close()
-
-	// Open a new database pointed to the same bucket storage
-	c, err = New(
-		WithLogger(logger),
-		WithReadWriteStorage(sinksource),
-	)
-	require.NoError(t, err)
-	defer c.Close()
-
-	// connect to our test db
-	db, err = c.DB(context.Background(), t.Name())
-	require.NoError(t, err)
-
-	// fetch new table
-	table, err = db.Table(t.Name(), config)
-	require.NoError(t, err)
-
-	bldr.Field(0).(*array.BinaryBuilder).Append([]byte("test"))
-	bldr.Field(1).(*array.BinaryBuilder).Append([]byte("value2"))
-	bldr.Field(2).(*array.BinaryBuilder).Append([]byte("value2"))
-
-	r = bldr.NewRecord()
-	defer r.Release()
-
-	_, err = table.InsertRecord(ctx, r)
-	require.NoError(t, err)
-}
-
 func Test_DB_Filter_Block(t *testing.T) {
 	sanitize := func(name string) string {
 		return strings.Replace(name, "/", "-", -1)
@@ -748,6 +504,9 @@ func Test_DB_Filter_Block(t *testing.T) {
 			db, err = c.DB(context.Background(), sanitize(t.Name()))
 			require.NoError(t, err)
 
+			_, err = db.Table(sanitize(t.Name()), config)
+			require.NoError(t, err)
+
 			engine := query.NewEngine(
 				memory.NewGoAllocator(),
 				db.TableProvider(),
@@ -1033,6 +792,9 @@ func Test_DB_Block_Optimization(t *testing.T) {
 			db, err = c.DB(context.Background(), sanitize(t.Name()))
 			require.NoError(t, err)
 
+			_, err = db.Table(sanitize(t.Name()), config)
+			require.NoError(t, err)
+
 			engine := query.NewEngine(
 				memory.NewGoAllocator(),
 				db.TableProvider(),
@@ -1946,6 +1708,9 @@ func Test_DB_EngineInMemory(t *testing.T) {
 	db, err = c.DB(context.Background(), "test")
 	require.NoError(t, err)
 
+	_, err = db.Table("test", config)
+	require.NoError(t, err)
+
 	pool := memory.NewCheckedAllocator(memory.DefaultAllocator)
 	defer pool.AssertSize(t, 0)
 	engine := query.NewEngine(pool, db.TableProvider(), query.WithPhysicalplanOptions(physicalplan.WithInMemoryOnly()))
diff --git a/query/logicalplan/builder_test.go b/query/logicalplan/builder_test.go
index ca50ac1..0ae0951 100644
--- a/query/logicalplan/builder_test.go
+++ b/query/logicalplan/builder_test.go
@@ -81,3 +81,21 @@ func TestLogicalPlanBuilderWithoutProjection(t *testing.T) {
 		},
 	}, p)
 }
+
+func TestRenamedColumn(t *testing.T) {
+	tableProvider := &mockTableProvider{schema: dynparquet.NewSampleSchema()}
+	_, err := (&Builder{}).
+		Scan(tableProvider, "table1").
+		Filter(Col("labels.test").Eq(Literal("abc"))).
+		Project(
+			Div(Mul(Col("value"), Literal(int64(2))), Literal(int64(2))).Alias("other_value"),
+			Col("stacktrace"),
+		).
+		Aggregate(
+			[]*AggregationFunction{Sum(Col("other_value"))},
+			[]Expr{Col("stacktrace")},
+		).
+		Project(Col("stacktrace"), Sum(Col("other_value")).Alias("value_sum")).
+		Build()
+	require.NoError(t, err)
+}
diff --git a/query/logicalplan/logicalplan_test.go b/query/logicalplan/logicalplan_test.go
index b151def..490dedf 100644
--- a/query/logicalplan/logicalplan_test.go
+++ b/query/logicalplan/logicalplan_test.go
@@ -56,7 +56,7 @@ func TestInputSchemaGetter(t *testing.T) {
 	schema := dynparquet.NewSampleSchema()
 
 	// test we can get the table by traversing to find the TableScan
-	plan, _ := (&Builder{}).
+	plan, err := (&Builder{}).
 		Scan(&mockTableProvider{schema}, "table1").
 		Filter(Col("labels.test").Eq(Literal("abc"))).
 		Aggregate(
@@ -65,31 +65,8 @@ func TestInputSchemaGetter(t *testing.T) {
 		).
 		Project(Col("stacktrace"), Sum(Col("value")).Alias("value_sum")).
 		Build()
+	require.NoError(t, err)
 	require.Equal(t, schema, plan.InputSchema())
-
-	// test we can get the table by traversing to find SchemaScan
-	plan, _ = (&Builder{}).
-		ScanSchema(&mockTableProvider{schema}, "table1").
-		Filter(Col("labels.test").Eq(Literal("abc"))).
-		Aggregate(
-			[]*AggregationFunction{Sum(Col("value"))},
-			[]Expr{Col("stacktrace")},
-		).
-		Project(Col("stacktrace"), Sum(Col("value")).Alias("value_sum")).
-		Build()
-	require.Equal(t, schema, plan.InputSchema())
-
-	// test it returns null in case where we built a logical plan w/ no
-	// TableScan or SchemaScan
-	plan, _ = (&Builder{}).
-		Filter(Col("labels.test").Eq(Literal("abc"))).
-		Aggregate(
-			[]*AggregationFunction{Sum(Col("value"))},
-			[]Expr{Col("stacktrace")},
-		).
-		Project(Col("stacktrace"), Sum(Col("value")).Alias("value_sum")).
-		Build()
-	require.Nil(t, plan.InputSchema())
 }
 
 func Test_ExprClone(t *testing.T) {
diff --git a/query/logicalplan/optimize_test.go b/query/logicalplan/optimize_test.go
index b369676..199123d 100644
--- a/query/logicalplan/optimize_test.go
+++ b/query/logicalplan/optimize_test.go
@@ -11,7 +11,7 @@ import (
 
 func TestOptimizePhysicalProjectionPushDown(t *testing.T) {
 	tableProvider := &mockTableProvider{schema: dynparquet.NewSampleSchema()}
-	p, _ := (&Builder{}).
+	p, err := (&Builder{}).
 		Scan(tableProvider, "table1").
 		Filter(Col("labels.test").Eq(Literal("abc"))).
 		Aggregate(
@@ -20,6 +20,7 @@ func TestOptimizePhysicalProjectionPushDown(t *testing.T) {
 		).
 		Project(Col("stacktrace"), Sum(Col("value")).Alias("value_sum")).
 		Build()
+	require.NoError(t, err)
 
 	optimizer := &PhysicalProjectionPushDown{}
 	optimizer.Optimize(p)
@@ -44,28 +45,26 @@ func TestOptimizePhysicalProjectionPushDown(t *testing.T) {
 }
 
 func TestOptimizeDistinctPushDown(t *testing.T) {
-	p, _ := (&Builder{}).
-		Scan(nil, "table1").
+	tableProvider := &mockTableProvider{schema: dynparquet.NewSampleSchema()}
+	p, err := (&Builder{}).
+		Scan(tableProvider, "table1").
 		Distinct(Col("labels.test")).
 		Build()
+	require.NoError(t, err)
 
 	optimizer := &DistinctPushDown{}
 	p = optimizer.Optimize(p)
 
-	require.Equal(t, &TableScan{
-		TableName: "table1",
-		Distinct: []Expr{
-			&Column{ColumnName: "labels.test"},
-		},
-	},
+	require.Equal(t,
+		[]Expr{&Column{ColumnName: "labels.test"}},
 		// Distinct -> TableScan
-		p.Input.Input.TableScan,
+		p.Input.Input.TableScan.Distinct,
 	)
 }
 
 func TestOptimizeFilterPushDown(t *testing.T) {
 	tableProvider := &mockTableProvider{schema: dynparquet.NewSampleSchema()}
-	p, _ := (&Builder{}).
+	p, err := (&Builder{}).
 		Scan(tableProvider, "table1").
 		Filter(Col("labels.test").Eq(Literal("abc"))).
 		Aggregate(
@@ -74,6 +73,7 @@ func TestOptimizeFilterPushDown(t *testing.T) {
 		).
 		Project(Col("stacktrace"), Sum(Col("value")).Alias("value_sum")).
 		Build()
+	require.NoError(t, err)
 
 	optimizer := &FilterPushDown{}
 	optimizer.Optimize(p)
@@ -96,7 +96,7 @@ func TestOptimizeFilterPushDown(t *testing.T) {
 }
 
 func TestProjectionPushDown(t *testing.T) {
-	p, _ := (&Builder{}).
+	p, err := (&Builder{}).
 		Scan(&mockTableProvider{schema: dynparquet.NewSampleSchema()}, "table1").
 		Filter(Col("labels.test").Eq(Literal("abc"))).
 		Aggregate(
@@ -104,6 +104,7 @@ func TestProjectionPushDown(t *testing.T) {
 			[]Expr{Col("stacktrace")},
 		).
 		Build()
+	require.NoError(t, err)
 
 	p = (&PhysicalProjectionPushDown{}).Optimize(p)
 
@@ -116,7 +117,7 @@ func TestProjectionPushDown(t *testing.T) {
 }
 
 func TestProjectionPushDownReset(t *testing.T) {
-	p, _ := (&Builder{}).
+	p, err := (&Builder{}).
 		Scan(&mockTableProvider{schema: dynparquet.NewSampleSchema()}, "table1").
 		Filter(Col("labels.test").Eq(Literal("abc"))).
 		Aggregate(
@@ -125,6 +126,7 @@ func TestProjectionPushDownReset(t *testing.T) {
 		).
 		Project(Col("test")).
 		Build()
+	require.NoError(t, err)
 
 	p = (&PhysicalProjectionPushDown{}).Optimize(p)
 
@@ -137,10 +139,11 @@ func TestProjectionPushDownReset(t *testing.T) {
 }
 
 func TestProjectionPushDownOfDistinct(t *testing.T) {
-	p, _ := (&Builder{}).
+	p, err := (&Builder{}).
 		Scan(&mockTableProvider{schema: dynparquet.NewSampleSchema()}, "table1").
 		Distinct(DynCol("labels")).
 		Build()
+	require.NoError(t, err)
 
 	p = (&PhysicalProjectionPushDown{}).Optimize(p)
 
@@ -149,7 +152,7 @@ func TestProjectionPushDownOfDistinct(t *testing.T) {
 
 func TestAllOptimizers(t *testing.T) {
 	tableProvider := &mockTableProvider{schema: dynparquet.NewSampleSchema()}
-	p, _ := (&Builder{}).
+	p, err := (&Builder{}).
 		Scan(tableProvider, "table1").
 		Filter(Col("labels.test").Eq(Literal("abc"))).
 		Aggregate(
@@ -158,6 +161,7 @@ func TestAllOptimizers(t *testing.T) {
 		).
 		Project(Col("stacktrace")).
 		Build()
+	require.NoError(t, err)
 
 	optimizers := DefaultOptimizers()
 
diff --git a/query/logicalplan/validate_test.go b/query/logicalplan/validate_test.go
index 46dd8d9..4fd50ca 100644
--- a/query/logicalplan/validate_test.go
+++ b/query/logicalplan/validate_test.go
@@ -82,7 +82,7 @@ func TestAggregationExprColumnMustExistInSchema(t *testing.T) {
 	require.True(t, strings.HasPrefix(planErr.message, "invalid aggregation"))
 	require.Len(t, planErr.children, 1)
 	exprErr := planErr.children[0]
-	require.True(t, strings.HasPrefix(exprErr.message, "column not found"))
+	require.Contains(t, exprErr.message, "column \"bad_column\" not found")
 }
 
 func TestAggregationCannotSumOrMaxTextColumn(t *testing.T) {
@@ -111,7 +111,7 @@ func TestAggregationCannotSumOrMaxTextColumn(t *testing.T) {
 		require.True(t, strings.HasPrefix(planErr.message, "invalid aggregation"))
 		require.Len(t, planErr.children, 1)
 		exprErr := planErr.children[0]
-		require.True(t, strings.HasPrefix(exprErr.message, testCase.errMsg))
+		require.Contains(t, exprErr.message, "expression type dictionary<values=binary, indices=uint32, ordered=false> is not supported")
 	}
 }