Fix code to use Expressions fields in VirtualTable instead of Literal

expr/expression.go

@@ -25,6 +25,13 @@ import (
 //
 //	NewScalarFunc(reg, id, nil, MustExpr(NewRootFieldRef(...)),
 //	    MustExpr(NewScalarFunc(...)))
+
+type (
+	// VirtualTableExpressionValue is a slice of other expression where each
+	// element in the slice is a different field in the struct
+	VirtualTableExpressionValue []Expression
+)
+
 func MustExpr(e Expression, err error) Expression {
 	if err != nil {
 		panic(err)
@@ -1603,3 +1610,33 @@ func (ex *Extended) ToProto() *proto.ExtendedExpression {
 		ReferredExpr:       refs,
 	}
 }
+
+func VirtualTableExpressionFromProto(s *proto.Expression_Nested_Struct, reg ExtensionRegistry) (VirtualTableExpressionValue, error) {
+	fields := make(VirtualTableExpressionValue, len(s.Fields))
+	for i, f := range s.Fields {
+		val, err := ExprFromProto(f, nil, reg)
+		if err != nil {
+			return nil, err
+		}
+		fields[i] = val
+	}
+	return fields, nil
+}
+
+func VirtualTableExprFromLiteralProto(s *proto.Expression_Literal_Struct) VirtualTableExpressionValue {
+	fields := make(VirtualTableExpressionValue, len(s.Fields))
+	for i, f := range s.Fields {
+		fields[i] = LiteralFromProto(f)
+	}
+	return fields
+}
+
+func (s VirtualTableExpressionValue) ToProto() *proto.Expression_Nested_Struct {
+	fields := make([]*proto.Expression, len(s))
+	for i, f := range s {
+		fields[i] = f.ToProto()
+	}
+	return &proto.Expression_Nested_Struct{
+		Fields: fields,
+	}
+}

expr/expressions_test.go

@@ -419,3 +419,43 @@ func TestRoundTripExtendedExpression(t *testing.T) {
 		assert.Truef(t, pb.Equal(&ex, out), "expected: %s\ngot: %s", &ex, out)
 	}
 }
+
+func TestVirtualTableExpressionFromProto(t *testing.T) {
+	// define extensions with no plan for now
+	const planExt = `{
+		"extensionUris": [
+			{
+				"extensionUriAnchor": 1,
+				"uri": "https://github.com/substrait-io/substrait/blob/main/extensions/functions_arithmetic.yaml"
+			}
+		],
+		"extensions": [
+			{
+				"extensionFunction": {
+					"extensionUriReference": 1,
+					"functionAnchor": 2,
+					"name": "add:i32_i32"
+				}
+			}
+		],
+		"relations": []
+	}`
+
+	var plan proto.Plan
+	if err := protojson.Unmarshal([]byte(planExt), &plan); err != nil {
+		panic(err)
+	}
+
+	// get the extension set
+	extSet := ext.GetExtensionSet(&plan)
+	literal1 := expr.NewPrimitiveLiteral(int32(1), false)
+	expr1 := literal1.ToProto()
+
+	reg := expr.NewExtensionRegistry(extSet, &ext.DefaultCollection)
+	rows := &proto.Expression_Nested_Struct{Fields: []*proto.Expression{
+		expr1,
+	}}
+	exprRows, err := expr.VirtualTableExpressionFromProto(rows, reg)
+	require.NoError(t, err)
+	require.Len(t, exprRows, 1)
+}

expr/literals.go

@@ -45,25 +45,6 @@ type (
 	Null types.Type
 )
 
-func StructLiteralFromProto(s *proto.Expression_Literal_Struct) StructLiteralValue {
-	fields := make(StructLiteralValue, len(s.Fields))
-	for i, f := range s.Fields {
-		fields[i] = LiteralFromProto(f)
-	}
-	return fields
-}
-
-func (s StructLiteralValue) ToProto() *proto.Expression_Literal_Struct {
-	fields := make([]*proto.Expression_Literal, len(s))
-	for i, f := range s {
-		fields[i] = f.ToProtoLiteral()
-	}
-
-	return &proto.Expression_Literal_Struct{
-		Fields: fields,
-	}
-}
-
 // Literal represents a specific literal of some type which could also
 // be a typed null or a nested type like a struct/map/list.
 //

plan/builders.go

@@ -482,6 +482,19 @@ func (b *builder) NamedScan(tableName []string, schema types.NamedStruct) *Named
 }
 
 func (b *builder) VirtualTableRemap(fieldNames []string, remap []int32, values ...expr.StructLiteralValue) (*VirtualTableReadRel, error) {
+	// convert Literal to Expression
+	exprs := make([]expr.VirtualTableExpressionValue, 0)
+	for _, row := range values {
+		rowExpr := make(expr.VirtualTableExpressionValue, 0)
+		for _, col := range row {
+			rowExpr = append(rowExpr, col)
+		}
+		exprs = append(exprs, rowExpr)
+	}
+	return b.VirtualTableFromExprRemap(fieldNames, remap, exprs...)
+}
+
+func (b *builder) VirtualTableFromExprRemap(fieldNames []string, remap []int32, values ...expr.VirtualTableExpressionValue) (*VirtualTableReadRel, error) {
 	if len(values) == 0 {
 		return nil, fmt.Errorf("%w: must provide at least one set of values for virtual table", substraitgo.ErrInvalidRel)
 	}

plan/plan.go

@@ -316,9 +316,19 @@ func RelFromProto(rel *proto.Rel, reg expr.ExtensionRegistry) (Rel, error) {
 				advExtension: readType.NamedTable.AdvancedExtension,
 			}
 		case *proto.ReadRel_VirtualTable_:
-			values := make([]expr.StructLiteralValue, len(readType.VirtualTable.Values))
-			for i, v := range readType.VirtualTable.Values {
-				values[i] = expr.StructLiteralFromProto(v)
+			if len(readType.VirtualTable.Values) > 0 && len(readType.VirtualTable.Expressions) > 0 {
+				return nil, fmt.Errorf("VirtualTable Value can't have both liternal and expression")
+			}
+			var values []expr.VirtualTableExpressionValue
+			for _, v := range readType.VirtualTable.Values {
+				values = append(values, expr.VirtualTableExprFromLiteralProto(v))
+			}
+			for _, v := range readType.VirtualTable.Expressions {
+				row, err := expr.VirtualTableExpressionFromProto(v, reg)
+				if err != nil {
+					return nil, err
+				}
+				values = append(values, row)
 			}
 
 			out = &VirtualTableReadRel{

plan/plan_builder_test.go

@@ -1348,17 +1348,17 @@ func TestSetRelations(t *testing.T) {
 											}
 										},
 										"virtualTable": {
-											"values": [
+											"expressions": [
 												{
 													"fields": [
-														{ "string": "foo", "nullable": false },
-														{ "fp32": 1.5, "nullable": false }
+														{"literal": { "string": "foo", "nullable": false }},
+														{"literal": { "fp32": 1.5, "nullable": false }}
 													]
 												},
 												{
 													"fields": [
-														{ "string": "bar", "nullable": false },
-														{ "fp32": 3.5, "nullable": false }
+														{"literal": { "string": "bar", "nullable": false }},
+														{"literal": { "fp32": 3.5, "nullable": false }}
 													]
 												}
 											]
@@ -1429,7 +1429,7 @@ func TestEmptyVirtualTable(t *testing.T) {
 								}
 							},
 							"virtualTable": {
-								"values": [
+								"expressions": [
 									{},
 									{},
 									{},

plan/plan_test.go

@@ -0,0 +1,52 @@
+package plan
+
+import (
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/stretchr/testify/require"
+	"github.com/substrait-io/substrait-go/expr"
+	"github.com/substrait-io/substrait-go/extensions"
+	"github.com/substrait-io/substrait-go/proto"
+	"google.golang.org/protobuf/testing/protocmp"
+)
+
+func TestRelFromProto(t *testing.T) {
+
+	registry := expr.NewEmptyExtensionRegistry(&extensions.DefaultCollection)
+	literal5 := &proto.Expression_Literal{LiteralType: &proto.Expression_Literal_I64{I64: 5}}
+	exprLiteral5 := &proto.Expression{RexType: &proto.Expression_Literal_{Literal: literal5}}
+
+	nestedStructExpr1 := &proto.Expression_Nested_Struct{Fields: []*proto.Expression{exprLiteral5}}
+	virtualTableWithExpression := &proto.ReadRel_VirtualTable_{VirtualTable: &proto.ReadRel_VirtualTable{Expressions: []*proto.Expression_Nested_Struct{nestedStructExpr1}}}
+	readRelWithExpression := &proto.ReadRel{ReadType: virtualTableWithExpression}
+
+	literalStruct := &proto.Expression_Literal_Struct{Fields: []*proto.Expression_Literal{literal5}}
+	virtualTableWithLiteral := &proto.ReadRel_VirtualTable_{VirtualTable: &proto.ReadRel_VirtualTable{Values: []*proto.Expression_Literal_Struct{literalStruct}}}
+	readRelWithLiteral := &proto.ReadRel{ReadType: virtualTableWithLiteral}
+
+	for _, td := range []struct {
+		name     string
+		readType *proto.ReadRel
+	}{
+		{"virtual table with expression", readRelWithExpression},
+		{"virtual table with deprecated literal", readRelWithLiteral},
+	} {
+		t.Run(td.name, func(t *testing.T) {
+			rel := &proto.Rel{RelType: &proto.Rel_Read{Read: td.readType}}
+
+			outRel, err := RelFromProto(rel, registry)
+			require.NoError(t, err)
+			gotRel := outRel.ToProto()
+			gotReadRel, ok := gotRel.RelType.(*proto.Rel_Read)
+			require.True(t, ok)
+			gotVirtualTableReadRel, ok := gotReadRel.Read.ReadType.(*proto.ReadRel_VirtualTable_)
+			require.True(t, ok)
+			// in case of both deprecated or new expression, the output should be the same as the new expression
+			if diff := cmp.Diff(gotVirtualTableReadRel, virtualTableWithExpression, protocmp.Transform()); diff != "" {
+				t.Errorf("expression proto didn't match, diff:\n%v", diff)
+			}
+		})
+	}
+
+}

plan/relations.go

@@ -208,23 +208,23 @@ func (n *NamedTableReadRel) CopyWithExpressionRewrite(rewriteFunc RewriteFunc, _
 type VirtualTableReadRel struct {
 	baseReadRel
 
-	values []expr.StructLiteralValue
+	values []expr.VirtualTableExpressionValue
 }
 
-func (v *VirtualTableReadRel) Values() []expr.StructLiteralValue {
+func (v *VirtualTableReadRel) Values() []expr.VirtualTableExpressionValue {
 	return v.values
 }
 
 func (v *VirtualTableReadRel) ToProto() *proto.Rel {
 	readRel := v.toReadRelProto()
-	values := make([]*proto.Expression_Literal_Struct, len(v.values))
+	values := make([]*proto.Expression_Nested_Struct, len(v.values))
 	for i, v := range v.values {
 		values[i] = v.ToProto()
 	}
 
 	readRel.ReadType = &proto.ReadRel_VirtualTable_{
 		VirtualTable: &proto.ReadRel_VirtualTable{
-			Values: values,
+			Expressions: values,
 		},
 	}
 	return &proto.Rel{

plan/relations_test.go

@@ -13,7 +13,7 @@ func noOpRewrite(e expr.Expression) (expr.Expression, error) {
 }
 
 func createVirtualTableReadRel(value int64) *VirtualTableReadRel {
-	return &VirtualTableReadRel{values: []expr.StructLiteralValue{[]expr.Literal{&expr.PrimitiveLiteral[int64]{Value: value}}}}
+	return &VirtualTableReadRel{values: []expr.VirtualTableExpressionValue{[]expr.Expression{&expr.PrimitiveLiteral[int64]{Value: value}}}}
 }
 
 func createPrimitiveFloat(value float64) expr.Expression {
@@ -40,7 +40,7 @@ func TestRelations_Copy(t *testing.T) {
 	projectRel := &ProjectRel{input: createVirtualTableReadRel(1), exprs: []expr.Expression{createPrimitiveFloat(1.0), createPrimitiveFloat(2.0)}}
 	setRel := &SetRel{inputs: []Rel{createVirtualTableReadRel(1), createVirtualTableReadRel(2), createVirtualTableReadRel(3)}, op: SetOpUnionAll}
 	sortRel := &SortRel{input: createVirtualTableReadRel(1), sorts: []expr.SortField{{Expr: createPrimitiveFloat(1.0), Kind: types.SortAscNullsFirst}}}
-	virtualTableReadRel := &VirtualTableReadRel{values: []expr.StructLiteralValue{[]expr.Literal{&expr.PrimitiveLiteral[int64]{Value: 1}}}}
+	virtualTableReadRel := &VirtualTableReadRel{values: []expr.VirtualTableExpressionValue{[]expr.Expression{&expr.PrimitiveLiteral[int64]{Value: 1}}}}
 	namedTableWriteRel := &NamedTableWriteRel{input: namedTableReadRel}
 
 	type relationTestCase struct {