feat: support subquery expressions (#134)

Author: Adam-Alani

README.md

@@ -18,7 +18,7 @@ This is work in progress still, things still to do:
   - [x] MultiOrList
   - [x] Cast
   - [x] Nested
-  - [ ] Subquery
+  - [x] Subquery
 - [ ] Serialization/Deserialization of Plan and Relations
   - [x] Plan
   - [x] PlanRel

expr/expression.go

@@ -317,6 +317,10 @@ func ExprFromProto(e *proto.Expression, baseSchema *types.RecordType, reg Extens
 	case *proto.Expression_Enum_:
 		return nil, fmt.Errorf("%w: deprecated", substraitgo.ErrNotImplemented)
 	case *proto.Expression_Subquery_:
+		if reg.subqueryConverter == nil {
+			return nil, fmt.Errorf("%w: subquery expressions require a subquery converter to be configured", substraitgo.ErrNotImplemented)
+		}
+		return reg.SubqueryFromProto(et.Subquery, baseSchema, reg)
 	}
 
 	return nil, fmt.Errorf("%w: ExprFromProto: %s", substraitgo.ErrNotImplemented, e)

expr/expressions_test.go

@@ -449,3 +449,172 @@ func TestCastVisit(t *testing.T) {
 		})
 	}
 }
+
+func TestSubqueryExpressionRoundtrip(t *testing.T) {
+	const substraitExtURI = "https://github.com/substrait-io/substrait/blob/main/extensions/functions_arithmetic.yaml"
+	// define extensions with no plan for now
+	const planExt = `{
+		"extensionUris": [
+			{
+				"extensionUriAnchor": 1,
+				"uri": "` + substraitExtURI + `"
+			}
+		],
+		"extensions": [],
+		"relations": []
+	}`
+
+	var planProto proto.Plan
+	if err := protojson.Unmarshal([]byte(planExt), &planProto); err != nil {
+		panic(err)
+	}
+
+	// get the extension set and create registry with subquery handler
+	extSet := ext.GetExtensionSet(&planProto)
+	c := ext.GetDefaultCollectionWithNoError()
+
+	// Create extension registry with subquery handler properly
+	baseReg := expr.NewExtensionRegistry(extSet, c)
+	subqueryReg := &plan.ExpressionConverter{ExtensionRegistry: baseReg}
+
+	// Create a simple mock relation for subqueries - single column of int32
+	mockSchema := types.NamedStruct{
+		Names: []string{"col1"},
+		Struct: types.StructType{
+			Types: []types.Type{&types.Int32Type{}},
+		},
+	}
+	mockRel := plan.NewBuilderDefault().NamedScan([]string{"test_table"}, mockSchema)
+
+	// Create base schema for needle expressions
+	baseSchema := types.NewRecordTypeFromTypes([]types.Type{&types.Int32Type{}, &types.StringType{}})
+
+	tests := []struct {
+		name    string
+		subExpr expr.Expression
+	}{
+		{
+			name:    "ScalarSubquery",
+			subExpr: plan.NewScalarSubquery(mockRel),
+		},
+		{
+			name: "InPredicateSubquery",
+			subExpr: plan.NewInPredicateSubquery(
+				[]expr.Expression{expr.NewPrimitiveLiteral(int32(42), false)},
+				mockRel,
+			),
+		},
+		{
+			name: "InPredicateSubquery_MultipleNeedles",
+			subExpr: func() expr.Expression {
+				// Create a 2-column relation for multi-needle test
+				twoColSchema := types.NamedStruct{
+					Names: []string{"col1", "col2"},
+					Struct: types.StructType{
+						Types: []types.Type{&types.Int32Type{}, &types.StringType{}},
+					},
+				}
+				twoColRel := plan.NewBuilderDefault().NamedScan([]string{"two_col_table"}, twoColSchema)
+
+				return plan.NewInPredicateSubquery(
+					[]expr.Expression{
+						expr.NewPrimitiveLiteral(int32(42), false),
+						expr.NewPrimitiveLiteral("test", false),
+					},
+					twoColRel,
+				)
+			}(),
+		},
+		{
+			name: "SetPredicateSubquery_EXISTS",
+			subExpr: plan.NewSetPredicateSubquery(
+				proto.Expression_Subquery_SetPredicate_PREDICATE_OP_EXISTS,
+				mockRel,
+			),
+		},
+		{
+			name: "SetPredicateSubquery_UNIQUE",
+			subExpr: plan.NewSetPredicateSubquery(
+				proto.Expression_Subquery_SetPredicate_PREDICATE_OP_UNIQUE,
+				mockRel,
+			),
+		},
+		{
+			name: "SetComparisonSubquery_ANY_EQ",
+			subExpr: plan.NewSetComparisonSubquery(
+				proto.Expression_Subquery_SetComparison_REDUCTION_OP_ANY,
+				proto.Expression_Subquery_SetComparison_COMPARISON_OP_EQ,
+				expr.NewPrimitiveLiteral(int32(42), false),
+				mockRel,
+			),
+		},
+		{
+			name: "SetComparisonSubquery_ALL_GT",
+			subExpr: plan.NewSetComparisonSubquery(
+				proto.Expression_Subquery_SetComparison_REDUCTION_OP_ALL,
+				proto.Expression_Subquery_SetComparison_COMPARISON_OP_GT,
+				expr.NewPrimitiveLiteral(int32(100), false),
+				mockRel,
+			),
+		},
+		{
+			name: "SetComparisonSubquery_ANY_NE",
+			subExpr: plan.NewSetComparisonSubquery(
+				proto.Expression_Subquery_SetComparison_REDUCTION_OP_ANY,
+				proto.Expression_Subquery_SetComparison_COMPARISON_OP_NE,
+				expr.NewPrimitiveLiteral(int32(0), false),
+				mockRel,
+			),
+		},
+		{
+			name: "SetComparisonSubquery_ALL_LE",
+			subExpr: plan.NewSetComparisonSubquery(
+				proto.Expression_Subquery_SetComparison_REDUCTION_OP_ALL,
+				proto.Expression_Subquery_SetComparison_COMPARISON_OP_LE,
+				expr.NewPrimitiveLiteral(int32(50), false),
+				mockRel,
+			),
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Convert expression to protobuf
+			protoExpr := tt.subExpr.ToProto()
+			require.NotNil(t, protoExpr)
+			require.NotNil(t, protoExpr.GetSubquery())
+
+			// Convert back from protobuf using ExprFromProto with subquery handler
+			baseReg.SetSubqueryConverter(subqueryReg)
+			fromProto, err := expr.ExprFromProto(protoExpr, baseSchema, baseReg)
+			require.NoError(t, err)
+			require.NotNil(t, fromProto)
+
+			// Verify that we got the right type of subquery back
+			switch tt.subExpr.(type) {
+			case *plan.ScalarSubquery:
+				assert.IsType(t, &plan.ScalarSubquery{}, fromProto)
+			case *plan.InPredicateSubquery:
+				assert.IsType(t, &plan.InPredicateSubquery{}, fromProto)
+			case *plan.SetPredicateSubquery:
+				assert.IsType(t, &plan.SetPredicateSubquery{}, fromProto)
+			case *plan.SetComparisonSubquery:
+				assert.IsType(t, &plan.SetComparisonSubquery{}, fromProto)
+			}
+
+			// Verify protobuf roundtrip
+			roundtripProto := fromProto.ToProto()
+			assert.True(t, pb.Equal(protoExpr, roundtripProto), "protobuf roundtrip failed")
+
+			// Verify basic properties
+			assert.Equal(t, tt.subExpr.IsScalar(), fromProto.IsScalar())
+			assert.True(t, tt.subExpr.GetType().Equals(fromProto.GetType()))
+
+			// Note: We don't test Equals() here because the current implementation
+			// of isRelEqual() only does pointer equality, so relations created from
+			// protobuf will never be equal to the original relations, even if they
+			// have identical content. This is a known limitation noted in the TODO
+			// comment in plan/subquery.go
+		})
+	}
+}

expr/field_reference.go

@@ -12,6 +12,13 @@ import (
 	"golang.org/x/exp/slices"
 )
 
+// RootRefType is a marker interface for types that can be used as a Root
+// reference in a FieldReference.
+//
+// A field reference is composed of two parts: a Root reference, which is the
+// output of an expression in this relation or a previous one, and a
+// ReferenceSegment or MaskedExpression, which allows referencing data within
+// that data structure - e.g. a field in a struct, or a value in a list or map.
 type RootRefType interface {
 	isRootRef()
 }

expr/utils.go

@@ -2,11 +2,41 @@
 
 package expr
 
-import "github.com/substrait-io/substrait-go/v4/extensions"
+import (
+	"github.com/substrait-io/substrait-go/v4/extensions"
+	"github.com/substrait-io/substrait-go/v4/types"
+	proto "github.com/substrait-io/substrait-protobuf/go/substraitpb"
+)
 
+// ExtensionRegistry provides functionality to resolve extension references and handle subquery expressions.
+// It combines an extensions.Set for looking up extension definitions with a Collection for extension metadata.
 type ExtensionRegistry struct {
 	extensions.Set
 	c *extensions.Collection
+
+	// subqueryConverter is injected by the plan package to handle subquery expressions
+	// TODO: We may want to consider refactoring to make a cleaner interface here
+	subqueryConverter
+}
+
+// subqueryConverter converts subqueries and the Relations within from the native
+// protobuf format into an Expression.
+//
+// This interface is private to avoid exposing the dependency cycle - a Subquery
+// contains a Plan, so the implementor of this has to exist in / import the plan
+// package, which we can't do here without creating a cycle with the expr
+// package.
+//
+// TODO: We may want to refactor this interface to be more generic or use a
+// different approach to avoid the cycle.
+type subqueryConverter interface {
+	SubqueryFromProto(sub *proto.Expression_Subquery, baseSchema *types.RecordType, reg ExtensionRegistry) (Expression, error)
+}
+
+// SetSubqueryConverter allows the plan package to inject a subquery converter.
+// This is an internal function used to break the dependency cycle between expr and plan packages.
+func (e *ExtensionRegistry) SetSubqueryConverter(converter subqueryConverter) {
+	e.subqueryConverter = converter
 }
 
 // NewExtensionRegistry creates a new registry.  If you have an existing plan you can use GetExtensionSet() to

plan/builders.go

@@ -149,6 +149,23 @@ type Builder interface {
 	// GetRelBuilder returns an expr.RelBuilder that can be used to construct
 	// relations which need multiple stages to build them.
 	GetRelBuilder() *RelBuilder
+
+	// Subquery expression builder methods
+
+	// InPredicateSubquery creates an IN predicate subquery expression that checks
+	// if the needles (left expressions) are contained in the haystack (right subquery).
+	InPredicateSubquery(needles []expr.Expression, haystack Rel) (*InPredicateSubquery, error)
+
+	// SetPredicateSubquery creates a set predicate subquery expression that checks
+	// if the subquery returns any rows.
+	SetPredicateSubquery(input Rel, predicateOp SetPredicateOp) (*SetPredicateSubquery, error)
+
+	// ScalarSubquery creates a scalar subquery expression that returns a single value.
+	ScalarSubquery(input Rel) (*ScalarSubquery, error)
+
+	// SetComparisonSubquery creates a set comparison subquery expression that checks
+	// if the left expression is contained in the right subquery.
+	SetComparisonSubquery(left expr.Expression, right Rel, reductionOp SetComparisonReductionOp, comparisonOp SetComparisonComparisonOp) (*SetComparisonSubquery, error)
 }
 
 const FETCH_COUNT_ALL_RECORDS = -1
@@ -870,3 +887,87 @@ func (arb *AggregateRelBuilder) validate() error {
 
 	return nil
 }
+
+func (b *builder) InPredicateSubquery(needles []expr.Expression, haystack Rel) (*InPredicateSubquery, error) {
+	if haystack == nil {
+		return nil, errNilInputRel
+	}
+
+	if len(needles) == 0 {
+		return nil, fmt.Errorf("%w: IN predicate subquery must have at least one needle expression",
+			substraitgo.ErrInvalidExpr)
+	}
+
+	for i, needle := range needles {
+		if needle == nil {
+			return nil, fmt.Errorf("%w: needle expression %d cannot be nil",
+				substraitgo.ErrInvalidExpr, i)
+		}
+	}
+
+	// Validate that the number of needle expressions matches the number of columns in the haystack
+	haystackSchema := haystack.RecordType()
+	if len(needles) != int(haystackSchema.FieldCount()) {
+		return nil, fmt.Errorf("%w: number of needle expressions (%d) must match number of columns in haystack (%d)",
+			substraitgo.ErrInvalidExpr, len(needles), haystackSchema.FieldCount())
+	}
+
+	return NewInPredicateSubquery(needles, haystack), nil
+}
+
+// SetPredicateSubquery creates a subquery that tests for the existence or uniqueness of rows
+// in the input relation.
+func (b *builder) SetPredicateSubquery(input Rel, predicateOp SetPredicateOp) (*SetPredicateSubquery, error) {
+	if input == nil {
+		return nil, errNilInputRel
+	}
+
+	if predicateOp == SetPredicateOpUnspecified {
+		return nil, fmt.Errorf("predicateOp must be specified")
+	}
+
+	return NewSetPredicateSubquery(
+		predicateOp,
+		input,
+	), nil
+}
+
+func (b *builder) ScalarSubquery(input Rel) (*ScalarSubquery, error) {
+	if input == nil {
+		return nil, errNilInputRel
+	}
+
+	return NewScalarSubquery(input), nil
+}
+
+// SetComparisonSubquery creates a subquery that compares a single expression against
+// a set of values from a relation using ANY or ALL operations with comparison operators.
+// The reductionOp determines whether to use ANY or ALL semantics, and the comparisonOp
+// specifies the comparison operator (e.g., =, !=, <, >, <=, >=).
+func (b *builder) SetComparisonSubquery(
+	left expr.Expression,
+	right Rel,
+	reductionOp SetComparisonReductionOp,
+	comparisonOp SetComparisonComparisonOp,
+) (*SetComparisonSubquery, error) {
+	if reductionOp == SetComparisonReductionOpUnspecified {
+		return nil, fmt.Errorf("reductionOp must be specified")
+	}
+	if comparisonOp == SetComparisonComparisonOpUnspecified {
+		return nil, fmt.Errorf("comparisonOp must be specified")
+	}
+
+	if left == nil {
+		return nil, errNilInputRel
+	}
+	if right == nil {
+		return nil, errNilInputRel
+	}
+
+	return NewSetComparisonSubquery(
+		reductionOp,
+		comparisonOp,
+		left,
+		right,
+	), nil
+}