Repository: polarsignals/frostdb
Pull request #746 (2024-02-19 16:55:51)

Add unique and "and" aggregations, and if and isNull projections

thorfour: May be worth adding some logic tests with these additions?

Patch:
diff --git a/query/logicalplan/expr.go b/query/logicalplan/expr.go
index 6f680ab..ca55c7f 100644
--- a/query/logicalplan/expr.go
+++ b/query/logicalplan/expr.go
@@ -724,6 +724,8 @@ const (
 	AggFuncMax
 	AggFuncCount
 	AggFuncAvg
+	AggFuncUnique
+	AggFuncAnd
 )
 
 func (f AggFunc) String() string {
@@ -738,6 +740,10 @@ func (f AggFunc) String() string {
 		return "count"
 	case AggFuncAvg:
 		return "avg"
+	case AggFuncUnique:
+		return "unique"
+	case AggFuncAnd:
+		return "and"
 	default:
 		panic("unknown aggregation function")
 	}
@@ -771,6 +777,20 @@ func Count(expr Expr) *AggregationFunction {
 	}
 }
 
+func Unique(expr Expr) *AggregationFunction {
+	return &AggregationFunction{
+		Func: AggFuncUnique,
+		Expr: expr,
+	}
+}
+
+func AndAgg(expr Expr) *AggregationFunction {
+	return &AggregationFunction{
+		Func: AggFuncAnd,
+		Expr: expr,
+	}
+}
+
 func Avg(expr Expr) *AggregationFunction {
 	return &AggregationFunction{
 		Func: AggFuncAvg,
@@ -778,6 +798,205 @@ func Avg(expr Expr) *AggregationFunction {
 	}
 }
 
+func IsNull(expr Expr) *IsNullExpr {
+	return &IsNullExpr{
+		Expr: expr,
+	}
+}
+
+type IsNullExpr struct {
+	Expr Expr
+}
+
+func (e *IsNullExpr) Equal(other Expr) bool {
+	if other == nil {
+		// if both are nil, they are equal
+		return e == nil
+	}
+
+	if isNull, ok := other.(*IsNullExpr); ok {
+		return e.Expr.Equal(isNull.Expr)
+	}
+
+	return false
+}
+
+func (e *IsNullExpr) Clone() Expr {
+	return &IsNullExpr{
+		Expr: e.Expr.Clone(),
+	}
+}
+
+func (e *IsNullExpr) DataType(l ExprTypeFinder) (arrow.DataType, error) {
+	_, err := e.Expr.DataType(l)
+	if err != nil {
+		return nil, err
+	}
+
+	return arrow.FixedWidthTypes.Boolean, nil
+}
+
+func (e *IsNullExpr) Accept(visitor Visitor) bool {
+	continu := visitor.PreVisit(e)
+	if !continu {
+		return false
+	}
+
+	continu = e.Expr.Accept(visitor)
+	if !continu {
+		return false
+	}
+
+	continu = visitor.Visit(e)
+	if !continu {
+		return false
+	}
+
+	return visitor.PostVisit(e)
+}
+
+func (e *IsNullExpr) Computed() bool {
+	return true
+}
+
+func (e *IsNullExpr) Name() string {
+	return "isnull(" + e.Expr.Name() + ")"
+}
+
+func (e *IsNullExpr) String() string { return e.Name() }
+
+func (e *IsNullExpr) ColumnsUsedExprs() []Expr {
+	return e.Expr.ColumnsUsedExprs()
+}
+
+func (e *IsNullExpr) MatchColumn(columnName string) bool {
+	return e.Name() == columnName
+}
+
+func (e *IsNullExpr) MatchPath(path string) bool {
+	return strings.HasPrefix(e.Name(), path)
+}
+
+func If(cond, then, els Expr) *IfExpr {
+	return &IfExpr{
+		Cond: cond,
+		Then: then,
+		Else: els,
+	}
+}
+
+type IfExpr struct {
+	Cond Expr
+	Then Expr
+	Else Expr
+}
+
+func (e *IfExpr) Equal(other Expr) bool {
+	if other == nil {
+		// if both are nil, they are equal
+		return e == nil
+	}
+
+	if ife, ok := other.(*IfExpr); ok {
+		return e.Cond.Equal(ife.Cond) && e.Then.Equal(ife.Then) && e.Else.Equal(ife.Else)
+	}
+
+	return false
+}
+
+func (e *IfExpr) Clone() Expr {
+	return &IfExpr{
+		Cond: e.Cond.Clone(),
+		Then: e.Then.Clone(),
+		Else: e.Else.Clone(),
+	}
+}
+
+func (e *IfExpr) DataType(l ExprTypeFinder) (arrow.DataType, error) {
+	condType, err := e.Cond.DataType(l)
+	if err != nil {
+		return nil, err
+	}
+
+	if !arrow.TypeEqual(condType, arrow.FixedWidthTypes.Boolean) {
+		return nil, fmt.Errorf("condition expression must be of type bool, got %s", condType)
+	}
+
+	thenType, err := e.Then.DataType(l)
+	if err != nil {
+		return nil, err
+	}
+
+	elseType, err := e.Else.DataType(l)
+	if err != nil {
+		return nil, err
+	}
+
+	if !arrow.TypeEqual(thenType, elseType) {
+		return nil, fmt.Errorf("then and else expression must have the same type, got %s and %s", thenType, elseType)
+	}
+
+	return thenType, nil
+}
+
+func (e *IfExpr) Accept(visitor Visitor) bool {
+	continu := visitor.PreVisit(e)
+	if !continu {
+		return false
+	}
+
+	continu = e.Cond.Accept(visitor)
+	if !continu {
+		return false
+	}
+
+	continu = e.Then.Accept(visitor)
+	if !continu {
+		return false
+	}
+
+	continu = e.Else.Accept(visitor)
+	if !continu {
+		return false
+	}
+
+	continu = visitor.Visit(e)
+	if !continu {
+		return false
+	}
+
+	return visitor.PostVisit(e)
+}
+
+func (e *IfExpr) Alias(name string) *AliasExpr {
+	return &AliasExpr{
+		Expr:  e,
+		Alias: name,
+	}
+}
+
+func (e *IfExpr) Computed() bool {
+	return true
+}
+
+func (e *IfExpr) Name() string {
+	return "if(" + e.Cond.Name() + ") { " + e.Then.Name() + " } else { " + e.Else.Name() + "}"
+}
+
+func (e *IfExpr) String() string { return e.Name() }
+
+func (e *IfExpr) ColumnsUsedExprs() []Expr {
+	return append(append(e.Cond.ColumnsUsedExprs(), e.Then.ColumnsUsedExprs()...), e.Else.ColumnsUsedExprs()...)
+}
+
+func (e *IfExpr) MatchColumn(columnName string) bool {
+	return e.Cond.MatchColumn(columnName) || e.Then.MatchColumn(columnName) || e.Else.MatchColumn(columnName)
+}
+
+func (e *IfExpr) MatchPath(path string) bool {
+	return e.Cond.MatchPath(path) || e.Then.MatchPath(path) || e.Else.MatchPath(path)
+}
+
 type AliasExpr struct {
 	Expr  Expr
 	Alias string
@@ -975,7 +1194,19 @@ func (n *NotExpr) Equal(other Expr) bool {
 	return false
 }
 
-func (n *NotExpr) DataType(ExprTypeFinder) (arrow.DataType, error) { return nil, nil }
+func (n *NotExpr) DataType(l ExprTypeFinder) (arrow.DataType, error) {
+	typ, err := n.Expr.DataType(l)
+	if err != nil {
+		return nil, err
+	}
+
+	if !arrow.TypeEqual(typ, arrow.FixedWidthTypes.Boolean) {
+		return nil, fmt.Errorf("not expression can only be applied to boolean expressions, got %s", typ)
+	}
+
+	return arrow.FixedWidthTypes.Boolean, nil
+}
+
 func (n *NotExpr) Accept(visitor Visitor) bool {
 	continu := visitor.PreVisit(n)
 	if !continu {
diff --git a/query/logicalplan/validate.go b/query/logicalplan/validate.go
index b490478..75bf2fc 100644
--- a/query/logicalplan/validate.go
+++ b/query/logicalplan/validate.go
@@ -268,13 +268,23 @@ func ValidateAggregationExpr(plan *LogicalPlan) *ExprValidationError {
 			}
 		}
 
-		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(),
+		switch expr.Func {
+		case AggFuncSum, AggFuncMin, AggFuncMax, AggFuncCount, AggFuncAvg, AggFuncUnique:
+			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(),
+				}
+			}
+		case AggFuncAnd:
+			if t != arrow.FixedWidthTypes.Boolean {
+				return &ExprValidationError{
+					expr:    expr.Expr,
+					message: fmt.Errorf("invalid aggregation: and aggregations can only aggregate bool type expressions, not %s", t).Error(),
+				}
 			}
 		}
 	}
diff --git a/query/physicalplan/aggregate.go b/query/physicalplan/aggregate.go
index 18fef41..f0feabe 100644
--- a/query/physicalplan/aggregate.go
+++ b/query/physicalplan/aggregate.go
@@ -91,6 +91,10 @@ func chooseAggregationFunction(
 		return &MaxAggregation{}, nil
 	case logicalplan.AggFuncCount:
 		return &CountAggregation{}, nil
+	case logicalplan.AggFuncUnique:
+		return &UniqueAggregation{}, nil
+	case logicalplan.AggFuncAnd:
+		return &AndAggregation{}, nil
 	default:
 		return nil, fmt.Errorf("unsupported aggregation function: %s", aggFunc.String())
 	}
@@ -609,7 +613,8 @@ func (a *HashAggregate) finishAggregate(ctx context.Context, aggIdx int, aggrega
 		groupByArrays = append(groupByArrays, aggregateArray)
 
 		aggregateFields = append(aggregateFields, arrow.Field{
-			Name: aggregation.resultName, Type: aggregateArray.DataType(),
+			Name: aggregation.resultName,
+			Type: aggregateArray.DataType(),
 		})
 	}
 
@@ -627,6 +632,105 @@ func (a *HashAggregate) finishAggregate(ctx context.Context, aggIdx int, aggrega
 	return nil
 }
 
+type AndAggregation struct{}
+
+var ErrUnsupportedAndType = errors.New("unsupported type for is and aggregation, expected bool")
+
+func (a *AndAggregation) Aggregate(pool memory.Allocator, arrs []arrow.Array) (arrow.Array, error) {
+	if len(arrs) == 0 {
+		return array.NewBooleanBuilder(pool).NewArray(), nil
+	}
+
+	typ := arrs[0].DataType().ID()
+	switch typ {
+	case arrow.BOOL:
+		return AndArrays(pool, arrs), nil
+	default:
+		return nil, fmt.Errorf("and array of %s: %w", typ, ErrUnsupportedAndType)
+	}
+}
+
+func AndArrays(pool memory.Allocator, arrs []arrow.Array) arrow.Array {
+	b := array.NewBooleanBuilder(pool)
+	defer b.Release()
+
+	for _, arr := range arrs {
+		if arr.Len() == 0 {
+			b.AppendNull()
+		}
+
+		arr := arr.(*array.Boolean)
+
+		val := true
+		for i := 0; i < arr.Len(); i++ {
+			if arr.IsValid(i) {
+				val = val && arr.Value(i)
+			}
+		}
+
+		b.Append(val)
+	}
+
+	return b.NewArray()
+}
+
+type UniqueAggregation struct{}
+
+var ErrUnsupportedIsUniqueType = errors.New("unsupported type for is unique aggregation, expected int64")
+
+func (a *UniqueAggregation) Aggregate(pool memory.Allocator, arrs []arrow.Array) (arrow.Array, error) {
+	if len(arrs) == 0 {
+		return array.NewInt64Builder(pool).NewArray(), nil
+	}
+
+	typ := arrs[0].DataType().ID()
+	switch typ {
+	case arrow.INT64:
+		return uniqueInt64arrays(pool, arrs), nil
+	default:
+		return nil, fmt.Errorf("isUnique array of %s: %w", typ, ErrUnsupportedIsUniqueType)
+	}
+}
+
+func uniqueInt64arrays(pool memory.Allocator, arrs []arrow.Array) arrow.Array {
+	res := array.NewInt64Builder(pool)
+	defer res.Release()
+
+	for _, arr := range arrs {
+		uniqueVal, isUnique, hasValues := int64ArrayHasUniqueValue(arr.(*array.Int64))
+		if !hasValues || !isUnique {
+			res.AppendNull()
+		} else {
+			res.Append(uniqueVal)
+		}
+	}
+
+	arr := res.NewArray()
+	return arr
+}
+
+func int64ArrayHasUniqueValue(arr *array.Int64) (int64, bool, bool) {
+	if arr.Len() == 0 {
+		return 0, false, false
+	}
+
+	if !arr.IsValid(0) {
+		return 0, false, true
+	}
+
+	val := arr.Value(0)
+	for i := 1; i < arr.Len(); i++ {
+		if !arr.IsValid(i) {
+			return 0, false, true
+		}
+		if val != arr.Value(i) {
+			return 0, false, true
+		}
+	}
+
+	return val, true, true
+}
+
 type SumAggregation struct{}
 
 var ErrUnsupportedSumType = errors.New("unsupported type for sum aggregation, expected int64 or float64")
diff --git a/query/physicalplan/project.go b/query/physicalplan/project.go
index 0162bce..6a518a9 100644
--- a/query/physicalplan/project.go
+++ b/query/physicalplan/project.go
@@ -485,6 +485,132 @@ func (p convertProjection) Project(mem memory.Allocator, ar arrow.Record) ([]arr
 	}}, []arrow.Array{c}, nil
 }
 
+type isNullProjection struct {
+	expr *logicalplan.IsNullExpr
+	p    columnProjection
+}
+
+func (p isNullProjection) Name() string {
+	return p.expr.String()
+}
+
+func (p isNullProjection) String() string {
+	return p.expr.String()
+}
+
+func (p isNullProjection) Project(mem memory.Allocator, ar arrow.Record) ([]arrow.Field, []arrow.Array, error) {
+	fields, cols, err := p.p.Project(mem, ar)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	if len(fields) != 1 || len(fields) != len(cols) {
+		return nil, nil, fmt.Errorf("invalid projection for isNull: expected 1 field and array, got %d fields %d arrays", len(fields), len(cols))
+	}
+
+	b := array.NewBooleanBuilder(mem)
+	defer b.Release()
+
+	b.Resize(cols[0].Len())
+
+	for i := 0; i < cols[0].Len(); i++ {
+		b.Append(cols[0].IsNull(i))
+	}
+
+	return []arrow.Field{{
+		Name:     p.expr.String(),
+		Type:     arrow.FixedWidthTypes.Boolean,
+		Nullable: fields[0].Nullable,
+		Metadata: fields[0].Metadata,
+	}}, []arrow.Array{b.NewBooleanArray()}, nil
+}
+
+type ifExprProjection struct {
+	expr *logicalplan.IfExpr
+
+	cond columnProjection
+	then columnProjection
+	els  columnProjection
+}
+
+func (p ifExprProjection) Name() string {
+	return p.expr.Name()
+}
+
+func (p ifExprProjection) String() string {
+	return p.expr.Name()
+}
+
+func (p ifExprProjection) Project(mem memory.Allocator, ar arrow.Record) ([]arrow.Field, []arrow.Array, error) {
+	f, thenCols, err := p.then.Project(mem, ar)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	_, elseCols, err := p.els.Project(mem, ar)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	_, condCols, err := p.cond.Project(mem, ar)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	if len(thenCols) != 1 || len(elseCols) != 1 || len(condCols) != 1 {
+		return nil, nil, fmt.Errorf("invalid projection for if: all columns must be non-empty and of equal length")
+	}
+
+	condCol := condCols[0]
+	thenCol := thenCols[0]
+	elseCol := elseCols[0]
+
+	condArr, ok := condCol.(*array.Boolean)
+	if !ok {
+		return nil, nil, fmt.Errorf("invalid projection for if: condition column must be of type boolean")
+	}
+
+	if condArr.Len() != thenCol.Len() || condArr.Len() != elseCol.Len() {
+		return nil, nil, fmt.Errorf("invalid projection for if: condition, then and else columns must be of equal length")
+	}
+
+	if !arrow.TypeEqual(thenCols[0].DataType(), elseCols[0].DataType()) {
+		return nil, nil, fmt.Errorf("invalid projection for if: then and else columns must be of the same type")
+	}
+
+	var res arrow.Array
+	switch thenCols[0].(type) {
+	case *array.Int64:
+		res = conditionalAddInt64(mem, condArr, thenCol.(*array.Int64), elseCol.(*array.Int64))
+	default:
+		return nil, nil, fmt.Errorf("unsupported if expression type: %s %T", thenCols[0].DataType(), thenCols[0])
+	}
+
+	return []arrow.Field{{
+		Name:     p.expr.Name(),
+		Type:     res.DataType(),
+		Nullable: f[0].Nullable,
+		Metadata: f[0].Metadata,
+	}}, []arrow.Array{res}, nil
+}
+
+func conditionalAddInt64(mem memory.Allocator, cond *array.Boolean, a, b *array.Int64) *array.Int64 {
+	res := array.NewInt64Builder(mem)
+	defer res.Release()
+
+	res.Resize(cond.Len())
+
+	for i := 0; i < cond.Len(); i++ {
+		if cond.IsValid(i) && cond.Value(i) {
+			res.Append(a.Value(i))
+		} else {
+			res.Append(b.Value(i))
+		}
+	}
+
+	return res.NewInt64Array()
+}
+
 type literalProjection struct {
 	value scalar.Scalar
 }
@@ -609,6 +735,38 @@ func projectionFromExpr(expr logicalplan.Expr) (columnProjection, error) {
 		default:
 			return nil, fmt.Errorf("unknown binary expression: %s", e.String())
 		}
+	case *logicalplan.IfExpr:
+		cond, err := projectionFromExpr(e.Cond)
+		if err != nil {
+			return nil, fmt.Errorf("condition projection for if projection: %w", err)
+		}
+
+		then, err := projectionFromExpr(e.Then)
+		if err != nil {
+			return nil, fmt.Errorf("then projection for if projection: %w", err)
+		}
+
+		els, err := projectionFromExpr(e.Else)
+		if err != nil {
+			return nil, fmt.Errorf("else projection for if projection: %w", err)
+		}
+
+		return ifExprProjection{
+			expr: e,
+			cond: cond,
+			then: then,
+			els:  els,
+		}, nil
+	case *logicalplan.IsNullExpr:
+		p, err := projectionFromExpr(e.Expr)
+		if err != nil {
+			return nil, fmt.Errorf("projection for is null projection: %w", err)
+		}
+
+		return isNullProjection{
+			expr: e,
+			p:    p,
+		}, nil
 	default:
 		return nil, fmt.Errorf("unsupported expression type for projection: %T", expr)
 	}


Test patch:
diff --git a/query/engine_test.go b/query/engine_test.go
new file mode 100644
index 0000000..83977ae
--- /dev/null
+++ b/query/engine_test.go
@@ -0,0 +1,219 @@
+package query
+
+import (
+	"context"
+	"testing"
+
+	"github.com/apache/arrow/go/v15/arrow"
+	"github.com/apache/arrow/go/v15/arrow/array"
+	"github.com/apache/arrow/go/v15/arrow/memory"
+	"github.com/stretchr/testify/require"
+
+	"github.com/polarsignals/frostdb/dynparquet"
+	schemapb "github.com/polarsignals/frostdb/gen/proto/go/frostdb/schema/v1alpha1"
+	"github.com/polarsignals/frostdb/query/logicalplan"
+)
+
+func TestUniqueAggregation(t *testing.T) {
+	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+	schema, err := dynparquet.SchemaFromDefinition(&schemapb.Schema{
+		Name: "test",
+		Columns: []*schemapb.Column{{
+			Name: "example",
+			StorageLayout: &schemapb.StorageLayout{
+				Type: schemapb.StorageLayout_TYPE_INT64,
+			},
+		}, {
+			Name: "timestamp",
+			StorageLayout: &schemapb.StorageLayout{
+				Type: schemapb.StorageLayout_TYPE_INT64,
+			},
+		}},
+	})
+	require.NoError(t, err)
+
+	rb := array.NewRecordBuilder(mem, arrow.NewSchema([]arrow.Field{{
+		Name: "example",
+		Type: arrow.PrimitiveTypes.Int64,
+	}, {
+		Name: "timestamp",
+		Type: arrow.PrimitiveTypes.Int64,
+	}}, nil))
+	defer rb.Release()
+
+	rb.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 2, 3}, nil)
+	rb.Field(1).(*array.Int64Builder).AppendValues([]int64{1, 1, 3}, nil)
+
+	ran := false
+	err = NewEngine(mem, &FakeTableProvider{
+		Tables: map[string]logicalplan.TableReader{
+			"test": &FakeTableReader{
+				FrostdbSchema: schema,
+				Records:       []arrow.Record{rb.NewRecord()},
+			},
+		},
+	}).ScanTable("test").
+		Aggregate(
+			[]*logicalplan.AggregationFunction{logicalplan.Unique(logicalplan.Col("example"))},
+			[]logicalplan.Expr{logicalplan.Col("timestamp")},
+		).
+		Execute(context.Background(), func(ctx context.Context, r arrow.Record) error {
+			require.Equal(t, []int64{1, 3}, r.Column(0).(*array.Int64).Int64Values())
+			require.True(t, r.Column(1).(*array.Int64).IsNull(0))
+			require.True(t, r.Column(1).(*array.Int64).IsValid(1))
+			require.Equal(t, int64(3), r.Column(1).(*array.Int64).Value(1))
+			ran = true
+			return nil
+		})
+	require.NoError(t, err)
+	require.True(t, ran)
+}
+
+func TestAndAggregation(t *testing.T) {
+	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+	schema, err := dynparquet.SchemaFromDefinition(&schemapb.Schema{
+		Name: "test",
+		Columns: []*schemapb.Column{{
+			Name: "example",
+			StorageLayout: &schemapb.StorageLayout{
+				Type: schemapb.StorageLayout_TYPE_BOOL,
+			},
+		}, {
+			Name: "timestamp",
+			StorageLayout: &schemapb.StorageLayout{
+				Type: schemapb.StorageLayout_TYPE_INT64,
+			},
+		}},
+	})
+	require.NoError(t, err)
+
+	rb := array.NewRecordBuilder(mem, arrow.NewSchema([]arrow.Field{{
+		Name: "example",
+		Type: arrow.FixedWidthTypes.Boolean,
+	}, {
+		Name: "timestamp",
+		Type: arrow.PrimitiveTypes.Int64,
+	}}, nil))
+	defer rb.Release()
+
+	rb.Field(0).(*array.BooleanBuilder).AppendValues([]bool{true, false, true, true}, nil)
+	rb.Field(1).(*array.Int64Builder).AppendValues([]int64{1, 1, 3, 3}, nil)
+
+	ran := false
+	err = NewEngine(mem, &FakeTableProvider{
+		Tables: map[string]logicalplan.TableReader{
+			"test": &FakeTableReader{
+				FrostdbSchema: schema,
+				Records:       []arrow.Record{rb.NewRecord()},
+			},
+		},
+	}).ScanTable("test").
+		Aggregate(
+			[]*logicalplan.AggregationFunction{logicalplan.AndAgg(logicalplan.Col("example"))},
+			[]logicalplan.Expr{logicalplan.Col("timestamp")},
+		).
+		Execute(context.Background(), func(ctx context.Context, r arrow.Record) error {
+			require.Equal(t, []int64{1, 3}, r.Column(0).(*array.Int64).Int64Values())
+			require.False(t, r.Column(1).(*array.Boolean).Value(0))
+			require.True(t, r.Column(1).(*array.Boolean).Value(1))
+			ran = true
+			return nil
+		})
+	require.NoError(t, err)
+	require.True(t, ran)
+}
+
+func TestIfProjection(t *testing.T) {
+	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+	schema, err := dynparquet.SchemaFromDefinition(&schemapb.Schema{
+		Name: "test",
+		Columns: []*schemapb.Column{{
+			Name: "example",
+			StorageLayout: &schemapb.StorageLayout{
+				Type: schemapb.StorageLayout_TYPE_INT64,
+			},
+		}, {
+			Name: "timestamp",
+			StorageLayout: &schemapb.StorageLayout{
+				Type: schemapb.StorageLayout_TYPE_INT64,
+			},
+		}},
+	})
+	require.NoError(t, err)
+
+	rb := array.NewRecordBuilder(mem, arrow.NewSchema([]arrow.Field{{
+		Name: "example",
+		Type: arrow.PrimitiveTypes.Int64,
+	}, {
+		Name: "timestamp",
+		Type: arrow.PrimitiveTypes.Int64,
+	}}, nil))
+	defer rb.Release()
+
+	rb.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 2, 3}, nil)
+	rb.Field(1).(*array.Int64Builder).AppendValues([]int64{1, 1, 3}, nil)
+
+	ran := false
+	err = NewEngine(mem, &FakeTableProvider{
+		Tables: map[string]logicalplan.TableReader{
+			"test": &FakeTableReader{
+				FrostdbSchema: schema,
+				Records:       []arrow.Record{rb.NewRecord()},
+			},
+		},
+	}).ScanTable("test").
+		Project(
+			logicalplan.If(logicalplan.Col("example").Eq(logicalplan.Literal(int64(1))), logicalplan.Literal(int64(1)), logicalplan.Literal(int64(0))),
+		).
+		Execute(context.Background(), func(ctx context.Context, r arrow.Record) error {
+			require.Equal(t, []int64{1, 0, 0}, r.Column(0).(*array.Int64).Int64Values())
+			ran = true
+			return nil
+		})
+	require.NoError(t, err)
+	require.True(t, ran)
+}
+
+func TestIsNullProjection(t *testing.T) {
+	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+	schema, err := dynparquet.SchemaFromDefinition(&schemapb.Schema{
+		Name: "test",
+		Columns: []*schemapb.Column{{
+			Name: "example",
+			StorageLayout: &schemapb.StorageLayout{
+				Type: schemapb.StorageLayout_TYPE_INT64,
+			},
+		}},
+	})
+	require.NoError(t, err)
+
+	rb := array.NewRecordBuilder(mem, arrow.NewSchema([]arrow.Field{{
+		Name: "example",
+		Type: arrow.PrimitiveTypes.Int64,
+	}}, nil))
+	defer rb.Release()
+
+	rb.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 2, 3}, []bool{true, true, false})
+
+	ran := false
+	err = NewEngine(mem, &FakeTableProvider{
+		Tables: map[string]logicalplan.TableReader{
+			"test": &FakeTableReader{
+				FrostdbSchema: schema,
+				Records:       []arrow.Record{rb.NewRecord()},
+			},
+		},
+	}).ScanTable("test").
+		Project(
+			logicalplan.IsNull(logicalplan.Col("example")),
+		).
+		Execute(context.Background(), func(ctx context.Context, r arrow.Record) error {
+			require.False(t, r.Column(0).(*array.Boolean).Value(0))
+			require.False(t, r.Column(0).(*array.Boolean).Value(1))
+			require.True(t, r.Column(0).(*array.Boolean).Value(2))
+			ran = true
+			return nil
+		})
+	require.NoError(t, err)
+	require.True(t, ran)
+}
diff --git a/query/testing.go b/query/testing.go
new file mode 100644
index 0000000..72b31d8
--- /dev/null
+++ b/query/testing.go
@@ -0,0 +1,69 @@
+package query
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"github.com/apache/arrow/go/v15/arrow"
+	"github.com/apache/arrow/go/v15/arrow/memory"
+
+	"github.com/polarsignals/frostdb/dynparquet"
+	"github.com/polarsignals/frostdb/query/logicalplan"
+)
+
+type FakeTableReader struct {
+	Records       []arrow.Record
+	FrostdbSchema *dynparquet.Schema
+}
+
+func (r *FakeTableReader) View(ctx context.Context, fn func(ctx context.Context, tx uint64) error) error {
+	return fn(ctx, 0)
+}
+
+func (r *FakeTableReader) Iterator(
+	ctx context.Context,
+	_ uint64,
+	_ memory.Allocator,
+	callbacks []logicalplan.Callback,
+	_ ...logicalplan.Option,
+) error {
+	if len(callbacks) == 0 {
+		return errors.New("no callbacks provided")
+	}
+
+	for i, r := range r.Records {
+		cb := callbacks[i%len(callbacks)]
+		if err := cb(ctx, r); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (r *FakeTableReader) SchemaIterator(
+	_ context.Context,
+	_ uint64,
+	_ memory.Allocator,
+	_ []logicalplan.Callback,
+	_ ...logicalplan.Option,
+) error {
+	return errors.New("not implemented")
+}
+
+func (r *FakeTableReader) Schema() *dynparquet.Schema {
+	return r.FrostdbSchema
+}
+
+type FakeTableProvider struct {
+	Tables map[string]logicalplan.TableReader
+}
+
+func (f *FakeTableProvider) GetTable(name string) (logicalplan.TableReader, error) {
+	if t, ok := f.Tables[name]; ok {
+		return t, nil
+	}
+
+	return nil, fmt.Errorf("table %s not found", name)
+}